screen_capture.rs

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