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