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