1use crate::{
2 DevicePixels, ForegroundExecutor, Size,
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,
11};
12use core_foundation::base::TCFType;
13use core_graphics::display::{
14 CGDirectDisplayID, CGDisplayCopyDisplayMode, CGDisplayModeGetPixelHeight,
15 CGDisplayModeGetPixelWidth, CGDisplayModeRelease,
16};
17use ctor::ctor;
18use futures::channel::oneshot;
19use media::core_media::{CMSampleBuffer, CMSampleBufferRef};
20use metal::NSInteger;
21use objc::{
22 class,
23 declare::ClassDecl,
24 msg_send,
25 runtime::{Class, Object, Sel},
26 sel, sel_impl,
27};
28use std::{cell::RefCell, ffi::c_void, mem, ptr, rc::Rc};
29
30#[derive(Clone)]
31pub struct MacScreenCaptureSource {
32 sc_display: id,
33}
34
35pub struct MacScreenCaptureStream {
36 sc_stream: id,
37 sc_stream_output: id,
38}
39
40static mut DELEGATE_CLASS: *const Class = ptr::null();
41static mut OUTPUT_CLASS: *const Class = ptr::null();
42const FRAME_CALLBACK_IVAR: &str = "frame_callback";
43
44#[allow(non_upper_case_globals)]
45const SCStreamOutputTypeScreen: NSInteger = 0;
46
47impl ScreenCaptureSource for MacScreenCaptureSource {
48 fn resolution(&self) -> Result<Size<DevicePixels>> {
49 unsafe {
50 let display_id: CGDirectDisplayID = msg_send![self.sc_display, displayID];
51 let display_mode_ref = CGDisplayCopyDisplayMode(display_id);
52 let width = CGDisplayModeGetPixelWidth(display_mode_ref);
53 let height = CGDisplayModeGetPixelHeight(display_mode_ref);
54 CGDisplayModeRelease(display_mode_ref);
55
56 Ok(size(
57 DevicePixels(width as i32),
58 DevicePixels(height as i32),
59 ))
60 }
61 }
62
63 fn stream(
64 &self,
65 _foreground_executor: &ForegroundExecutor,
66 frame_callback: Box<dyn Fn(ScreenCaptureFrame) + Send>,
67 ) -> oneshot::Receiver<Result<Box<dyn ScreenCaptureStream>>> {
68 unsafe {
69 let stream: id = msg_send![class!(SCStream), alloc];
70 let filter: id = msg_send![class!(SCContentFilter), alloc];
71 let configuration: id = msg_send![class!(SCStreamConfiguration), alloc];
72 let delegate: id = msg_send![DELEGATE_CLASS, alloc];
73 let output: id = msg_send![OUTPUT_CLASS, alloc];
74
75 let excluded_windows = NSArray::array(nil);
76 let filter: id = msg_send![filter, initWithDisplay:self.sc_display excludingWindows:excluded_windows];
77 let configuration: id = msg_send![configuration, init];
78 let _: id = msg_send![configuration, setScalesToFit: true];
79 let _: id = msg_send![configuration, setPixelFormat: 0x42475241];
80 // let _: id = msg_send![configuration, setShowsCursor: false];
81 // let _: id = msg_send![configuration, setCaptureResolution: 3];
82 let delegate: id = msg_send![delegate, init];
83 let output: id = msg_send![output, init];
84
85 output.as_mut().unwrap().set_ivar(
86 FRAME_CALLBACK_IVAR,
87 Box::into_raw(Box::new(frame_callback)) as *mut c_void,
88 );
89
90 let resolution = self.resolution().unwrap();
91 let _: id = msg_send![configuration, setWidth: resolution.width.0 as i64];
92 let _: id = msg_send![configuration, setHeight: resolution.height.0 as i64];
93 let stream: id = msg_send![stream, initWithFilter:filter configuration:configuration delegate:delegate];
94
95 let (mut tx, rx) = oneshot::channel();
96
97 let mut error: id = nil;
98 let _: () = msg_send![stream, addStreamOutput:output type:SCStreamOutputTypeScreen sampleHandlerQueue:0 error:&mut error as *mut id];
99 if error != nil {
100 let message: id = msg_send![error, localizedDescription];
101 tx.send(Err(anyhow!("failed to add stream output {message:?}")))
102 .ok();
103 return rx;
104 }
105
106 let tx = Rc::new(RefCell::new(Some(tx)));
107 let handler = ConcreteBlock::new({
108 move |error: id| {
109 let result = if error == nil {
110 let stream = MacScreenCaptureStream {
111 sc_stream: stream,
112 sc_stream_output: output,
113 };
114 Ok(Box::new(stream) as Box<dyn ScreenCaptureStream>)
115 } else {
116 let message: id = msg_send![error, localizedDescription];
117 Err(anyhow!("failed to stop screen capture stream {message:?}"))
118 };
119 if let Some(tx) = tx.borrow_mut().take() {
120 tx.send(result).ok();
121 }
122 }
123 });
124 let handler = handler.copy();
125 let _: () = msg_send![stream, startCaptureWithCompletionHandler:handler];
126 rx
127 }
128 }
129}
130
131impl Drop for MacScreenCaptureSource {
132 fn drop(&mut self) {
133 unsafe {
134 let _: () = msg_send![self.sc_display, release];
135 }
136 }
137}
138
139impl ScreenCaptureStream for MacScreenCaptureStream {}
140
141impl Drop for MacScreenCaptureStream {
142 fn drop(&mut self) {
143 unsafe {
144 let mut error: id = nil;
145 let _: () = msg_send![self.sc_stream, removeStreamOutput:self.sc_stream_output type:SCStreamOutputTypeScreen error:&mut error as *mut _];
146 if error != nil {
147 let message: id = msg_send![error, localizedDescription];
148 log::error!("failed to add stream output {message:?}");
149 }
150
151 let handler = ConcreteBlock::new(move |error: id| {
152 if error != nil {
153 let message: id = msg_send![error, localizedDescription];
154 log::error!("failed to stop screen capture stream {message:?}");
155 }
156 });
157 let block = handler.copy();
158 let _: () = msg_send![self.sc_stream, stopCaptureWithCompletionHandler:block];
159 let _: () = msg_send![self.sc_stream, release];
160 let _: () = msg_send![self.sc_stream_output, release];
161 }
162 }
163}
164
165pub(crate) fn get_sources() -> oneshot::Receiver<Result<Vec<Box<dyn ScreenCaptureSource>>>> {
166 unsafe {
167 let (mut tx, rx) = oneshot::channel();
168 let tx = Rc::new(RefCell::new(Some(tx)));
169
170 let block = ConcreteBlock::new(move |shareable_content: id, error: id| {
171 let Some(mut tx) = tx.borrow_mut().take() else {
172 return;
173 };
174 let result = if error == nil {
175 let displays: id = msg_send![shareable_content, displays];
176 let mut result = Vec::new();
177 for i in 0..displays.count() {
178 let display = displays.objectAtIndex(i);
179 let source = MacScreenCaptureSource {
180 sc_display: msg_send![display, retain],
181 };
182 result.push(Box::new(source) as Box<dyn ScreenCaptureSource>);
183 }
184 Ok(result)
185 } else {
186 let msg: id = msg_send![error, localizedDescription];
187 Err(anyhow!("Failed to register: {:?}", msg))
188 };
189 tx.send(result).ok();
190 });
191 let block = block.copy();
192
193 let _: () = msg_send![
194 class!(SCShareableContent),
195 getShareableContentExcludingDesktopWindows:YES
196 onScreenWindowsOnly:YES
197 completionHandler:block];
198 rx
199 }
200}
201
202#[ctor]
203unsafe fn build_classes() {
204 let mut decl = ClassDecl::new("GPUIStreamDelegate", class!(NSObject)).unwrap();
205 unsafe {
206 decl.add_method(
207 sel!(outputVideoEffectDidStartForStream:),
208 output_video_effect_did_start_for_stream as extern "C" fn(&Object, Sel, id),
209 );
210 decl.add_method(
211 sel!(outputVideoEffectDidStopForStream:),
212 output_video_effect_did_stop_for_stream as extern "C" fn(&Object, Sel, id),
213 );
214 decl.add_method(
215 sel!(stream:didStopWithError:),
216 stream_did_stop_with_error as extern "C" fn(&Object, Sel, id, id),
217 );
218 DELEGATE_CLASS = decl.register();
219
220 let mut decl = ClassDecl::new("GPUIStreamOutput", class!(NSObject)).unwrap();
221 decl.add_method(
222 sel!(stream:didOutputSampleBuffer:ofType:),
223 stream_did_output_sample_buffer_of_type
224 as extern "C" fn(&Object, Sel, id, id, NSInteger),
225 );
226 decl.add_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR);
227
228 OUTPUT_CLASS = decl.register();
229 }
230}
231
232extern "C" fn output_video_effect_did_start_for_stream(_this: &Object, _: Sel, _stream: id) {}
233
234extern "C" fn output_video_effect_did_stop_for_stream(_this: &Object, _: Sel, _stream: id) {}
235
236extern "C" fn stream_did_stop_with_error(_this: &Object, _: Sel, _stream: id, _error: id) {}
237
238extern "C" fn stream_did_output_sample_buffer_of_type(
239 this: &Object,
240 _: Sel,
241 _stream: id,
242 sample_buffer: id,
243 buffer_type: NSInteger,
244) {
245 if buffer_type != SCStreamOutputTypeScreen {
246 return;
247 }
248
249 unsafe {
250 let sample_buffer = sample_buffer as CMSampleBufferRef;
251 let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer);
252 if let Some(buffer) = sample_buffer.image_buffer() {
253 let callback: Box<Box<dyn Fn(ScreenCaptureFrame)>> =
254 Box::from_raw(*this.get_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR) as *mut _);
255 callback(ScreenCaptureFrame(buffer));
256 mem::forget(callback);
257 }
258 }
259}