screen_capture.rs

  1use crate::{
  2    DevicePixels, ForegroundExecutor, SharedString, SourceMetadata,
  3    platform::{ScreenCaptureFrame, ScreenCaptureSource, ScreenCaptureStream},
  4    size,
  5};
  6use anyhow::{Result, anyhow};
  7use block::ConcreteBlock;
  8use cocoa::{
  9    base::{YES, id, nil},
 10    foundation::{NSArray, NSString},
 11};
 12use collections::HashMap;
 13use core_foundation::base::TCFType;
 14use core_graphics::display::{
 15    CGDirectDisplayID, CGDisplayCopyDisplayMode, CGDisplayModeGetPixelHeight,
 16    CGDisplayModeGetPixelWidth, CGDisplayModeRelease,
 17};
 18use ctor::ctor;
 19use futures::channel::oneshot;
 20use media::core_media::{CMSampleBuffer, CMSampleBufferRef};
 21use metal::NSInteger;
 22use objc::{
 23    class,
 24    declare::ClassDecl,
 25    msg_send,
 26    runtime::{Class, Object, Sel},
 27    sel, sel_impl,
 28};
 29use std::{cell::RefCell, ffi::c_void, mem, ptr, rc::Rc};
 30
 31use super::NSStringExt;
 32
 33#[derive(Clone)]
 34pub struct MacScreenCaptureSource {
 35    sc_display: id,
 36    meta: Option<ScreenMeta>,
 37}
 38
 39pub struct MacScreenCaptureStream {
 40    sc_stream: id,
 41    sc_stream_output: id,
 42    meta: SourceMetadata,
 43}
 44
 45static mut DELEGATE_CLASS: *const Class = ptr::null();
 46static mut OUTPUT_CLASS: *const Class = ptr::null();
 47const FRAME_CALLBACK_IVAR: &str = "frame_callback";
 48
 49#[allow(non_upper_case_globals)]
 50const SCStreamOutputTypeScreen: NSInteger = 0;
 51
 52impl ScreenCaptureSource for MacScreenCaptureSource {
 53    fn metadata(&self) -> Result<SourceMetadata> {
 54        let (display_id, size) = unsafe {
 55            let display_id: CGDirectDisplayID = msg_send![self.sc_display, displayID];
 56            let display_mode_ref = CGDisplayCopyDisplayMode(display_id);
 57            let width = CGDisplayModeGetPixelWidth(display_mode_ref);
 58            let height = CGDisplayModeGetPixelHeight(display_mode_ref);
 59            CGDisplayModeRelease(display_mode_ref);
 60
 61            (
 62                display_id,
 63                size(DevicePixels(width as i32), DevicePixels(height as i32)),
 64            )
 65        };
 66        let (label, is_main) = self
 67            .meta
 68            .clone()
 69            .map(|meta| (meta.label, meta.is_main))
 70            .unzip();
 71
 72        Ok(SourceMetadata {
 73            id: display_id as u64,
 74            label,
 75            is_main,
 76            resolution: size,
 77        })
 78    }
 79
 80    fn stream(
 81        &self,
 82        _foreground_executor: &ForegroundExecutor,
 83        frame_callback: Box<dyn Fn(ScreenCaptureFrame) + Send>,
 84    ) -> oneshot::Receiver<Result<Box<dyn ScreenCaptureStream>>> {
 85        unsafe {
 86            let stream: id = msg_send![class!(SCStream), alloc];
 87            let filter: id = msg_send![class!(SCContentFilter), alloc];
 88            let configuration: id = msg_send![class!(SCStreamConfiguration), alloc];
 89            let delegate: id = msg_send![DELEGATE_CLASS, alloc];
 90            let output: id = msg_send![OUTPUT_CLASS, alloc];
 91
 92            let excluded_windows = NSArray::array(nil);
 93            let filter: id = msg_send![filter, initWithDisplay:self.sc_display excludingWindows:excluded_windows];
 94            let configuration: id = msg_send![configuration, init];
 95            let _: id = msg_send![configuration, setScalesToFit: true];
 96            let _: id = msg_send![configuration, setPixelFormat: 0x42475241];
 97            // let _: id = msg_send![configuration, setShowsCursor: false];
 98            // let _: id = msg_send![configuration, setCaptureResolution: 3];
 99            let delegate: id = msg_send![delegate, init];
100            let output: id = msg_send![output, init];
101
102            output.as_mut().unwrap().set_ivar(
103                FRAME_CALLBACK_IVAR,
104                Box::into_raw(Box::new(frame_callback)) as *mut c_void,
105            );
106
107            let meta = self.metadata().unwrap();
108            let _: id = msg_send![configuration, setWidth: meta.resolution.width.0 as i64];
109            let _: id = msg_send![configuration, setHeight: meta.resolution.height.0 as i64];
110            let stream: id = msg_send![stream, initWithFilter:filter configuration:configuration delegate:delegate];
111
112            let (mut tx, rx) = oneshot::channel();
113
114            let mut error: id = nil;
115            let _: () = msg_send![stream, addStreamOutput:output type:SCStreamOutputTypeScreen sampleHandlerQueue:0 error:&mut error as *mut id];
116            if error != nil {
117                let message: id = msg_send![error, localizedDescription];
118                tx.send(Err(anyhow!("failed to add stream  output {message:?}")))
119                    .ok();
120                return rx;
121            }
122
123            let tx = Rc::new(RefCell::new(Some(tx)));
124            let handler = ConcreteBlock::new({
125                move |error: id| {
126                    let result = if error == nil {
127                        let stream = MacScreenCaptureStream {
128                            meta: meta.clone(),
129                            sc_stream: stream,
130                            sc_stream_output: output,
131                        };
132                        Ok(Box::new(stream) as Box<dyn ScreenCaptureStream>)
133                    } else {
134                        let message: id = msg_send![error, localizedDescription];
135                        Err(anyhow!("failed to stop screen capture stream {message:?}"))
136                    };
137                    if let Some(tx) = tx.borrow_mut().take() {
138                        tx.send(result).ok();
139                    }
140                }
141            });
142            let handler = handler.copy();
143            let _: () = msg_send![stream, startCaptureWithCompletionHandler:handler];
144            rx
145        }
146    }
147}
148
149impl Drop for MacScreenCaptureSource {
150    fn drop(&mut self) {
151        unsafe {
152            let _: () = msg_send![self.sc_display, release];
153        }
154    }
155}
156
157impl ScreenCaptureStream for MacScreenCaptureStream {
158    fn metadata(&self) -> Result<SourceMetadata> {
159        Ok(self.meta.clone())
160    }
161}
162
163impl Drop for MacScreenCaptureStream {
164    fn drop(&mut self) {
165        unsafe {
166            let mut error: id = nil;
167            let _: () = msg_send![self.sc_stream, removeStreamOutput:self.sc_stream_output type:SCStreamOutputTypeScreen error:&mut error as *mut _];
168            if error != nil {
169                let message: id = msg_send![error, localizedDescription];
170                log::error!("failed to add stream  output {message:?}");
171            }
172
173            let handler = ConcreteBlock::new(move |error: id| {
174                if error != nil {
175                    let message: id = msg_send![error, localizedDescription];
176                    log::error!("failed to stop screen capture stream {message:?}");
177                }
178            });
179            let block = handler.copy();
180            let _: () = msg_send![self.sc_stream, stopCaptureWithCompletionHandler:block];
181            let _: () = msg_send![self.sc_stream, release];
182            let _: () = msg_send![self.sc_stream_output, release];
183        }
184    }
185}
186
187#[derive(Clone)]
188struct ScreenMeta {
189    label: SharedString,
190    // Is this the screen with menu bar?
191    is_main: bool,
192}
193
194unsafe fn screen_id_to_human_label() -> HashMap<CGDirectDisplayID, ScreenMeta> {
195    let screens: id = msg_send![class!(NSScreen), screens];
196    let count: usize = msg_send![screens, count];
197    let mut map = HashMap::default();
198    let screen_number_key = unsafe { NSString::alloc(nil).init_str("NSScreenNumber") };
199    for i in 0..count {
200        let screen: id = msg_send![screens, objectAtIndex: i];
201        let device_desc: id = msg_send![screen, deviceDescription];
202        if device_desc == nil {
203            continue;
204        }
205
206        let nsnumber: id = msg_send![device_desc, objectForKey: screen_number_key];
207        if nsnumber == nil {
208            continue;
209        }
210
211        let screen_id: u32 = msg_send![nsnumber, unsignedIntValue];
212
213        let name: id = msg_send![screen, localizedName];
214        if name != nil {
215            let cstr: *const std::os::raw::c_char = msg_send![name, UTF8String];
216            let rust_str = unsafe {
217                std::ffi::CStr::from_ptr(cstr)
218                    .to_string_lossy()
219                    .into_owned()
220            };
221            map.insert(
222                screen_id,
223                ScreenMeta {
224                    label: rust_str.into(),
225                    is_main: i == 0,
226                },
227            );
228        }
229    }
230    map
231}
232
233pub(crate) fn get_sources() -> oneshot::Receiver<Result<Vec<Rc<dyn ScreenCaptureSource>>>> {
234    unsafe {
235        let (mut tx, rx) = oneshot::channel();
236        let tx = Rc::new(RefCell::new(Some(tx)));
237        let screen_id_to_label = screen_id_to_human_label();
238        let block = ConcreteBlock::new(move |shareable_content: id, error: id| {
239            let Some(mut tx) = tx.borrow_mut().take() else {
240                return;
241            };
242
243            let result = if error == nil {
244                let displays: id = msg_send![shareable_content, displays];
245                let mut result = Vec::new();
246                for i in 0..displays.count() {
247                    let display = displays.objectAtIndex(i);
248                    let id: CGDirectDisplayID = msg_send![display, displayID];
249                    let meta = screen_id_to_label.get(&id).cloned();
250                    let source = MacScreenCaptureSource {
251                        sc_display: msg_send![display, retain],
252                        meta,
253                    };
254                    result.push(Rc::new(source) as Rc<dyn ScreenCaptureSource>);
255                }
256                Ok(result)
257            } else {
258                let msg: id = msg_send![error, localizedDescription];
259                Err(anyhow!(
260                    "Screen share failed: {:?}",
261                    NSStringExt::to_str(&msg)
262                ))
263            };
264            tx.send(result).ok();
265        });
266        let block = block.copy();
267
268        let _: () = msg_send![
269            class!(SCShareableContent),
270            getShareableContentExcludingDesktopWindows:YES
271                                   onScreenWindowsOnly:YES
272                                     completionHandler:block];
273        rx
274    }
275}
276
277#[ctor]
278unsafe fn build_classes() {
279    let mut decl = ClassDecl::new("GPUIStreamDelegate", class!(NSObject)).unwrap();
280    unsafe {
281        decl.add_method(
282            sel!(outputVideoEffectDidStartForStream:),
283            output_video_effect_did_start_for_stream as extern "C" fn(&Object, Sel, id),
284        );
285        decl.add_method(
286            sel!(outputVideoEffectDidStopForStream:),
287            output_video_effect_did_stop_for_stream as extern "C" fn(&Object, Sel, id),
288        );
289        decl.add_method(
290            sel!(stream:didStopWithError:),
291            stream_did_stop_with_error as extern "C" fn(&Object, Sel, id, id),
292        );
293        DELEGATE_CLASS = decl.register();
294
295        let mut decl = ClassDecl::new("GPUIStreamOutput", class!(NSObject)).unwrap();
296        decl.add_method(
297            sel!(stream:didOutputSampleBuffer:ofType:),
298            stream_did_output_sample_buffer_of_type
299                as extern "C" fn(&Object, Sel, id, id, NSInteger),
300        );
301        decl.add_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR);
302
303        OUTPUT_CLASS = decl.register();
304    }
305}
306
307const extern "C" fn output_video_effect_did_start_for_stream(_this: &Object, _: Sel, _stream: id) {}
308
309const extern "C" fn output_video_effect_did_stop_for_stream(_this: &Object, _: Sel, _stream: id) {}
310
311const extern "C" fn stream_did_stop_with_error(_this: &Object, _: Sel, _stream: id, _error: id) {}
312
313extern "C" fn stream_did_output_sample_buffer_of_type(
314    this: &Object,
315    _: Sel,
316    _stream: id,
317    sample_buffer: id,
318    buffer_type: NSInteger,
319) {
320    if buffer_type != SCStreamOutputTypeScreen {
321        return;
322    }
323
324    unsafe {
325        let sample_buffer = sample_buffer as CMSampleBufferRef;
326        let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer);
327        if let Some(buffer) = sample_buffer.image_buffer() {
328            let callback: Box<Box<dyn Fn(ScreenCaptureFrame)>> =
329                Box::from_raw(*this.get_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR) as *mut _);
330            callback(ScreenCaptureFrame(buffer));
331            mem::forget(callback);
332        }
333    }
334}