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            // Stream contains filter, configuration, and delegate internally so we release them here
114            // to prevent a memory leak when steam is dropped
115            let _: () = msg_send![filter, release];
116            let _: () = msg_send![configuration, release];
117            let _: () = msg_send![delegate, release];
118
119            let (mut tx, rx) = oneshot::channel();
120
121            let mut error: id = nil;
122            let _: () = msg_send![stream, addStreamOutput:output type:SCStreamOutputTypeScreen sampleHandlerQueue:0 error:&mut error as *mut id];
123            if error != nil {
124                let message: id = msg_send![error, localizedDescription];
125                let _: () = msg_send![stream, release];
126                let _: () = msg_send![output, release];
127                tx.send(Err(anyhow!("failed to add stream output {message:?}")))
128                    .ok();
129                return rx;
130            }
131
132            let tx = Rc::new(RefCell::new(Some(tx)));
133            let handler = ConcreteBlock::new({
134                move |error: id| {
135                    let result = if error == nil {
136                        let stream = MacScreenCaptureStream {
137                            meta: meta.clone(),
138                            sc_stream: stream,
139                            sc_stream_output: output,
140                        };
141                        Ok(Box::new(stream) as Box<dyn ScreenCaptureStream>)
142                    } else {
143                        let _: () = msg_send![stream, release];
144                        let _: () = msg_send![output, release];
145                        let message: id = msg_send![error, localizedDescription];
146                        Err(anyhow!("failed to start screen capture stream {message:?}"))
147                    };
148                    if let Some(tx) = tx.borrow_mut().take() {
149                        tx.send(result).ok();
150                    }
151                }
152            });
153            let handler = handler.copy();
154            let _: () = msg_send![stream, startCaptureWithCompletionHandler:handler];
155            rx
156        }
157    }
158}
159
160impl Drop for MacScreenCaptureSource {
161    fn drop(&mut self) {
162        unsafe {
163            let _: () = msg_send![self.sc_display, release];
164        }
165    }
166}
167
168impl ScreenCaptureStream for MacScreenCaptureStream {
169    fn metadata(&self) -> Result<SourceMetadata> {
170        Ok(self.meta.clone())
171    }
172}
173
174impl Drop for MacScreenCaptureStream {
175    fn drop(&mut self) {
176        unsafe {
177            let mut error: id = nil;
178            let _: () = msg_send![self.sc_stream, removeStreamOutput:self.sc_stream_output type:SCStreamOutputTypeScreen error:&mut error as *mut _];
179            if error != nil {
180                let message: id = msg_send![error, localizedDescription];
181                log::error!("failed to add stream  output {message:?}");
182            }
183
184            let handler = ConcreteBlock::new(move |error: id| {
185                if error != nil {
186                    let message: id = msg_send![error, localizedDescription];
187                    log::error!("failed to stop screen capture stream {message:?}");
188                }
189            });
190            let block = handler.copy();
191            let _: () = msg_send![self.sc_stream, stopCaptureWithCompletionHandler:block];
192            let _: () = msg_send![self.sc_stream, release];
193            let _: () = msg_send![self.sc_stream_output, release];
194        }
195    }
196}
197
198#[derive(Clone)]
199struct ScreenMeta {
200    label: SharedString,
201    // Is this the screen with menu bar?
202    is_main: bool,
203}
204
205unsafe fn screen_id_to_human_label() -> HashMap<CGDirectDisplayID, ScreenMeta> {
206    let screens: id = msg_send![class!(NSScreen), screens];
207    let count: usize = msg_send![screens, count];
208    let mut map = HashMap::default();
209    let screen_number_key = unsafe { ns_string("NSScreenNumber") };
210    for i in 0..count {
211        let screen: id = msg_send![screens, objectAtIndex: i];
212        let device_desc: id = msg_send![screen, deviceDescription];
213        if device_desc == nil {
214            continue;
215        }
216
217        let nsnumber: id = msg_send![device_desc, objectForKey: screen_number_key];
218        if nsnumber == nil {
219            continue;
220        }
221
222        let screen_id: u32 = msg_send![nsnumber, unsignedIntValue];
223
224        let name: id = msg_send![screen, localizedName];
225        if name != nil {
226            let cstr: *const std::os::raw::c_char = msg_send![name, UTF8String];
227            let rust_str = unsafe {
228                std::ffi::CStr::from_ptr(cstr)
229                    .to_string_lossy()
230                    .into_owned()
231            };
232            map.insert(
233                screen_id,
234                ScreenMeta {
235                    label: rust_str.into(),
236                    is_main: i == 0,
237                },
238            );
239        }
240    }
241    map
242}
243
244pub(crate) fn get_sources() -> oneshot::Receiver<Result<Vec<Rc<dyn ScreenCaptureSource>>>> {
245    unsafe {
246        let (mut tx, rx) = oneshot::channel();
247        let tx = Rc::new(RefCell::new(Some(tx)));
248        let screen_id_to_label = screen_id_to_human_label();
249        let block = ConcreteBlock::new(move |shareable_content: id, error: id| {
250            let Some(mut tx) = tx.borrow_mut().take() else {
251                return;
252            };
253
254            let result = if error == nil {
255                let displays: id = msg_send![shareable_content, displays];
256                let mut result = Vec::new();
257                for i in 0..displays.count() {
258                    let display = displays.objectAtIndex(i);
259                    let id: CGDirectDisplayID = msg_send![display, displayID];
260                    let meta = screen_id_to_label.get(&id).cloned();
261                    let source = MacScreenCaptureSource {
262                        sc_display: msg_send![display, retain],
263                        meta,
264                    };
265                    result.push(Rc::new(source) as Rc<dyn ScreenCaptureSource>);
266                }
267                Ok(result)
268            } else {
269                let msg: id = msg_send![error, localizedDescription];
270                Err(anyhow!(
271                    "Screen share failed: {:?}",
272                    NSStringExt::to_str(&msg)
273                ))
274            };
275            tx.send(result).ok();
276        });
277        let block = block.copy();
278
279        let _: () = msg_send![
280            class!(SCShareableContent),
281            getShareableContentExcludingDesktopWindows:YES
282                                   onScreenWindowsOnly:YES
283                                     completionHandler:block];
284        rx
285    }
286}
287
288#[ctor]
289unsafe fn build_classes() {
290    let mut decl = ClassDecl::new("GPUIStreamDelegate", class!(NSObject)).unwrap();
291    unsafe {
292        decl.add_method(
293            sel!(outputVideoEffectDidStartForStream:),
294            output_video_effect_did_start_for_stream as extern "C" fn(&Object, Sel, id),
295        );
296        decl.add_method(
297            sel!(outputVideoEffectDidStopForStream:),
298            output_video_effect_did_stop_for_stream as extern "C" fn(&Object, Sel, id),
299        );
300        decl.add_method(
301            sel!(stream:didStopWithError:),
302            stream_did_stop_with_error as extern "C" fn(&Object, Sel, id, id),
303        );
304        DELEGATE_CLASS = decl.register();
305
306        let mut decl = ClassDecl::new("GPUIStreamOutput", class!(NSObject)).unwrap();
307        decl.add_method(
308            sel!(stream:didOutputSampleBuffer:ofType:),
309            stream_did_output_sample_buffer_of_type
310                as extern "C" fn(&Object, Sel, id, id, NSInteger),
311        );
312        decl.add_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR);
313
314        OUTPUT_CLASS = decl.register();
315    }
316}
317
318extern "C" fn output_video_effect_did_start_for_stream(_this: &Object, _: Sel, _stream: id) {}
319
320extern "C" fn output_video_effect_did_stop_for_stream(_this: &Object, _: Sel, _stream: id) {}
321
322extern "C" fn stream_did_stop_with_error(_this: &Object, _: Sel, _stream: id, _error: id) {}
323
324extern "C" fn stream_did_output_sample_buffer_of_type(
325    this: &Object,
326    _: Sel,
327    _stream: id,
328    sample_buffer: id,
329    buffer_type: NSInteger,
330) {
331    if buffer_type != SCStreamOutputTypeScreen {
332        return;
333    }
334
335    unsafe {
336        let sample_buffer = sample_buffer as CMSampleBufferRef;
337        let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer);
338        if let Some(buffer) = sample_buffer.image_buffer() {
339            let callback: Box<Box<dyn Fn(ScreenCaptureFrame)>> =
340                Box::from_raw(*this.get_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR) as *mut _);
341            callback(ScreenCaptureFrame(buffer));
342            mem::forget(callback);
343        }
344    }
345}