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