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