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