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")]
40extern "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 decl.add_method(
204 sel!(outputVideoEffectDidStartForStream:),
205 output_video_effect_did_start_for_stream as extern "C" fn(&Object, Sel, id),
206 );
207 decl.add_method(
208 sel!(outputVideoEffectDidStopForStream:),
209 output_video_effect_did_stop_for_stream as extern "C" fn(&Object, Sel, id),
210 );
211 decl.add_method(
212 sel!(stream:didStopWithError:),
213 stream_did_stop_with_error as extern "C" fn(&Object, Sel, id, id),
214 );
215 DELEGATE_CLASS = decl.register();
216
217 let mut decl = ClassDecl::new("GPUIStreamOutput", class!(NSObject)).unwrap();
218 decl.add_method(
219 sel!(stream:didOutputSampleBuffer:ofType:),
220 stream_did_output_sample_buffer_of_type as extern "C" fn(&Object, Sel, id, id, NSInteger),
221 );
222 decl.add_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR);
223
224 OUTPUT_CLASS = decl.register();
225}
226
227extern "C" fn output_video_effect_did_start_for_stream(_this: &Object, _: Sel, _stream: id) {}
228
229extern "C" fn output_video_effect_did_stop_for_stream(_this: &Object, _: Sel, _stream: id) {}
230
231extern "C" fn stream_did_stop_with_error(_this: &Object, _: Sel, _stream: id, _error: id) {}
232
233extern "C" fn stream_did_output_sample_buffer_of_type(
234 this: &Object,
235 _: Sel,
236 _stream: id,
237 sample_buffer: id,
238 buffer_type: NSInteger,
239) {
240 if buffer_type != SCStreamOutputTypeScreen {
241 return;
242 }
243
244 unsafe {
245 let sample_buffer = sample_buffer as CMSampleBufferRef;
246 let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer);
247 if let Some(buffer) = sample_buffer.image_buffer() {
248 let callback: Box<Box<dyn Fn(ScreenCaptureFrame)>> =
249 Box::from_raw(*this.get_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR) as *mut _);
250 callback(ScreenCaptureFrame(buffer));
251 mem::forget(callback);
252 }
253 }
254}