1use super::ns_string;
2use crate::{
3 DevicePixels, ForegroundExecutor, SharedString, SourceMetadata,
4 platform::{ScreenCaptureFrame, ScreenCaptureSource, ScreenCaptureStream},
5 size,
6};
7use anyhow::{Result, anyhow};
8use block::ConcreteBlock;
9use cocoa::{
10 base::{NO, YES, id, nil},
11 foundation::NSArray,
12};
13use collections::HashMap;
14use core_foundation::base::TCFType;
15use core_graphics::{
16 base::CGFloat,
17 color_space::CGColorSpace,
18 display::{
19 CGDirectDisplayID, CGDisplayCopyDisplayMode, CGDisplayModeGetPixelHeight,
20 CGDisplayModeGetPixelWidth, CGDisplayModeRelease,
21 },
22 image::CGImage,
23};
24use core_video::pixel_buffer::CVPixelBuffer;
25use ctor::ctor;
26use foreign_types::ForeignType;
27use futures::channel::oneshot;
28use image::{ImageBuffer, Rgba, RgbaImage};
29use media::core_media::{CMSampleBuffer, CMSampleBufferRef};
30use metal::NSInteger;
31use objc::{
32 class,
33 declare::ClassDecl,
34 msg_send,
35 runtime::{Class, Object, Sel},
36 sel, sel_impl,
37};
38use std::{cell::RefCell, ffi::c_void, mem, ptr, rc::Rc};
39
40use super::NSStringExt;
41
42#[derive(Clone)]
43pub struct MacScreenCaptureSource {
44 sc_display: id,
45 meta: Option<ScreenMeta>,
46}
47
48pub struct MacScreenCaptureStream {
49 sc_stream: id,
50 sc_stream_output: id,
51 meta: SourceMetadata,
52}
53
54static mut DELEGATE_CLASS: *const Class = ptr::null();
55static mut OUTPUT_CLASS: *const Class = ptr::null();
56const FRAME_CALLBACK_IVAR: &str = "frame_callback";
57
58#[allow(non_upper_case_globals)]
59const SCStreamOutputTypeScreen: NSInteger = 0;
60
61impl ScreenCaptureSource for MacScreenCaptureSource {
62 fn metadata(&self) -> Result<SourceMetadata> {
63 let (display_id, size) = unsafe {
64 let display_id: CGDirectDisplayID = msg_send![self.sc_display, displayID];
65 let display_mode_ref = CGDisplayCopyDisplayMode(display_id);
66 let width = CGDisplayModeGetPixelWidth(display_mode_ref);
67 let height = CGDisplayModeGetPixelHeight(display_mode_ref);
68 CGDisplayModeRelease(display_mode_ref);
69
70 (
71 display_id,
72 size(DevicePixels(width as i32), DevicePixels(height as i32)),
73 )
74 };
75 let (label, is_main) = self
76 .meta
77 .clone()
78 .map(|meta| (meta.label, meta.is_main))
79 .unzip();
80
81 Ok(SourceMetadata {
82 id: display_id as u64,
83 label,
84 is_main,
85 resolution: size,
86 })
87 }
88
89 fn stream(
90 &self,
91 _foreground_executor: &ForegroundExecutor,
92 frame_callback: Box<dyn Fn(ScreenCaptureFrame) + Send>,
93 ) -> oneshot::Receiver<Result<Box<dyn ScreenCaptureStream>>> {
94 unsafe {
95 let stream: id = msg_send![class!(SCStream), alloc];
96 let filter: id = msg_send![class!(SCContentFilter), alloc];
97 let configuration: id = msg_send![class!(SCStreamConfiguration), alloc];
98 let delegate: id = msg_send![DELEGATE_CLASS, alloc];
99 let output: id = msg_send![OUTPUT_CLASS, alloc];
100
101 let excluded_windows = NSArray::array(nil);
102 let filter: id = msg_send![filter, initWithDisplay:self.sc_display excludingWindows:excluded_windows];
103 let configuration: id = msg_send![configuration, init];
104 let _: id = msg_send![configuration, setScalesToFit: true];
105 let _: id = msg_send![configuration, setPixelFormat: 0x42475241];
106 // let _: id = msg_send![configuration, setShowsCursor: false];
107 // let _: id = msg_send![configuration, setCaptureResolution: 3];
108 let delegate: id = msg_send![delegate, init];
109 let output: id = msg_send![output, init];
110
111 output.as_mut().unwrap().set_ivar(
112 FRAME_CALLBACK_IVAR,
113 Box::into_raw(Box::new(frame_callback)) as *mut c_void,
114 );
115
116 let meta = self.metadata().unwrap();
117 let _: id = msg_send![configuration, setWidth: meta.resolution.width.0 as i64];
118 let _: id = msg_send![configuration, setHeight: meta.resolution.height.0 as i64];
119 let stream: id = msg_send![stream, initWithFilter:filter configuration:configuration delegate:delegate];
120
121 let (mut tx, rx) = oneshot::channel();
122
123 let mut error: id = nil;
124 let _: () = msg_send![stream, addStreamOutput:output type:SCStreamOutputTypeScreen sampleHandlerQueue:0 error:&mut error as *mut id];
125 if error != nil {
126 let message: id = msg_send![error, localizedDescription];
127 tx.send(Err(anyhow!("failed to add stream output {message:?}")))
128 .ok();
129 return rx;
130 }
131
132 let tx = Rc::new(RefCell::new(Some(tx)));
133 let handler = ConcreteBlock::new({
134 move |error: id| {
135 let result = if error == nil {
136 let stream = MacScreenCaptureStream {
137 meta: meta.clone(),
138 sc_stream: stream,
139 sc_stream_output: output,
140 };
141 Ok(Box::new(stream) as Box<dyn ScreenCaptureStream>)
142 } else {
143 let message: id = msg_send![error, localizedDescription];
144 Err(anyhow!("failed to stop screen capture stream {message:?}"))
145 };
146 if let Some(tx) = tx.borrow_mut().take() {
147 tx.send(result).ok();
148 }
149 }
150 });
151 let handler = handler.copy();
152 let _: () = msg_send![stream, startCaptureWithCompletionHandler:handler];
153 rx
154 }
155 }
156}
157
158impl Drop for MacScreenCaptureSource {
159 fn drop(&mut self) {
160 unsafe {
161 let _: () = msg_send![self.sc_display, release];
162 }
163 }
164}
165
166impl ScreenCaptureStream for MacScreenCaptureStream {
167 fn metadata(&self) -> Result<SourceMetadata> {
168 Ok(self.meta.clone())
169 }
170}
171
172impl Drop for MacScreenCaptureStream {
173 fn drop(&mut self) {
174 unsafe {
175 let mut error: id = nil;
176 let _: () = msg_send![self.sc_stream, removeStreamOutput:self.sc_stream_output type:SCStreamOutputTypeScreen error:&mut error as *mut _];
177 if error != nil {
178 let message: id = msg_send![error, localizedDescription];
179 log::error!("failed to add stream output {message:?}");
180 }
181
182 let handler = ConcreteBlock::new(move |error: id| {
183 if error != nil {
184 let message: id = msg_send![error, localizedDescription];
185 log::error!("failed to stop screen capture stream {message:?}");
186 }
187 });
188 let block = handler.copy();
189 let _: () = msg_send![self.sc_stream, stopCaptureWithCompletionHandler:block];
190 let _: () = msg_send![self.sc_stream, release];
191 let _: () = msg_send![self.sc_stream_output, release];
192 }
193 }
194}
195
196#[derive(Clone)]
197struct ScreenMeta {
198 label: SharedString,
199 // Is this the screen with menu bar?
200 is_main: bool,
201}
202
203unsafe fn screen_id_to_human_label() -> HashMap<CGDirectDisplayID, ScreenMeta> {
204 let screens: id = msg_send![class!(NSScreen), screens];
205 let count: usize = msg_send![screens, count];
206 let mut map = HashMap::default();
207 let screen_number_key = unsafe { ns_string("NSScreenNumber") };
208 for i in 0..count {
209 let screen: id = msg_send![screens, objectAtIndex: i];
210 let device_desc: id = msg_send![screen, deviceDescription];
211 if device_desc == nil {
212 continue;
213 }
214
215 let nsnumber: id = msg_send![device_desc, objectForKey: screen_number_key];
216 if nsnumber == nil {
217 continue;
218 }
219
220 let screen_id: u32 = msg_send![nsnumber, unsignedIntValue];
221
222 let name: id = msg_send![screen, localizedName];
223 if name != nil {
224 let cstr: *const std::os::raw::c_char = msg_send![name, UTF8String];
225 let rust_str = unsafe {
226 std::ffi::CStr::from_ptr(cstr)
227 .to_string_lossy()
228 .into_owned()
229 };
230 map.insert(
231 screen_id,
232 ScreenMeta {
233 label: rust_str.into(),
234 is_main: i == 0,
235 },
236 );
237 }
238 }
239 map
240}
241
242pub(crate) fn get_sources() -> oneshot::Receiver<Result<Vec<Rc<dyn ScreenCaptureSource>>>> {
243 unsafe {
244 let (mut tx, rx) = oneshot::channel();
245 let tx = Rc::new(RefCell::new(Some(tx)));
246 let screen_id_to_label = screen_id_to_human_label();
247 let block = ConcreteBlock::new(move |shareable_content: id, error: id| {
248 let Some(mut tx) = tx.borrow_mut().take() else {
249 return;
250 };
251
252 let result = if error == nil {
253 let displays: id = msg_send![shareable_content, displays];
254 let mut result = Vec::new();
255 for i in 0..displays.count() {
256 let display = displays.objectAtIndex(i);
257 let id: CGDirectDisplayID = msg_send![display, displayID];
258 let meta = screen_id_to_label.get(&id).cloned();
259 let source = MacScreenCaptureSource {
260 sc_display: msg_send![display, retain],
261 meta,
262 };
263 result.push(Rc::new(source) as Rc<dyn ScreenCaptureSource>);
264 }
265 Ok(result)
266 } else {
267 let msg: id = msg_send![error, localizedDescription];
268 Err(anyhow!(
269 "Screen share failed: {:?}",
270 NSStringExt::to_str(&msg)
271 ))
272 };
273 tx.send(result).ok();
274 });
275 let block = block.copy();
276
277 let _: () = msg_send![
278 class!(SCShareableContent),
279 getShareableContentExcludingDesktopWindows:YES
280 onScreenWindowsOnly:YES
281 completionHandler:block];
282 rx
283 }
284}
285
286/// Captures a single screenshot of a specific window by its CGWindowID.
287///
288/// This uses ScreenCaptureKit's `initWithDesktopIndependentWindow:` API which can
289/// capture windows even when they are positioned off-screen (e.g., at -10000, -10000).
290///
291/// # Arguments
292/// * `window_id` - The CGWindowID (NSWindow's windowNumber) of the window to capture
293///
294/// # Returns
295/// An `RgbaImage` containing the captured window contents, or an error if capture failed.
296pub fn capture_window_screenshot(window_id: u32) -> oneshot::Receiver<Result<RgbaImage>> {
297 let (tx, rx) = oneshot::channel();
298 let tx = Rc::new(RefCell::new(Some(tx)));
299
300 unsafe {
301 log::info!(
302 "capture_window_screenshot: looking for window_id={}",
303 window_id
304 );
305 let content_handler = ConcreteBlock::new(move |shareable_content: id, error: id| {
306 log::info!("capture_window_screenshot: content handler called");
307 if error != nil {
308 if let Some(sender) = tx.borrow_mut().take() {
309 let msg: id = msg_send![error, localizedDescription];
310 sender
311 .send(Err(anyhow!(
312 "Failed to get shareable content: {:?}",
313 NSStringExt::to_str(&msg)
314 )))
315 .ok();
316 }
317 return;
318 }
319
320 let windows: id = msg_send![shareable_content, windows];
321 let count: usize = msg_send![windows, count];
322
323 let mut target_window: id = nil;
324 log::info!(
325 "capture_window_screenshot: searching {} windows for window_id={}",
326 count,
327 window_id
328 );
329 for i in 0..count {
330 let window: id = msg_send![windows, objectAtIndex: i];
331 let wid: u32 = msg_send![window, windowID];
332 if wid == window_id {
333 log::info!(
334 "capture_window_screenshot: found matching window at index {}",
335 i
336 );
337 target_window = window;
338 break;
339 }
340 }
341
342 if target_window == nil {
343 if let Some(sender) = tx.borrow_mut().take() {
344 sender
345 .send(Err(anyhow!(
346 "Window with ID {} not found in shareable content",
347 window_id
348 )))
349 .ok();
350 }
351 return;
352 }
353
354 log::info!("capture_window_screenshot: calling capture_window_frame");
355 capture_window_frame(target_window, &tx);
356 });
357 let content_handler = content_handler.copy();
358
359 let _: () = msg_send![
360 class!(SCShareableContent),
361 getShareableContentExcludingDesktopWindows:NO
362 onScreenWindowsOnly:NO
363 completionHandler:content_handler
364 ];
365 }
366
367 rx
368}
369
370unsafe fn capture_window_frame(
371 sc_window: id,
372 tx: &Rc<RefCell<Option<oneshot::Sender<Result<RgbaImage>>>>>,
373) {
374 log::info!("capture_window_frame: creating filter for window");
375 let filter: id = msg_send![class!(SCContentFilter), alloc];
376 let filter: id = msg_send![filter, initWithDesktopIndependentWindow: sc_window];
377 log::info!("capture_window_frame: filter created: {:?}", filter);
378
379 let configuration: id = msg_send![class!(SCStreamConfiguration), alloc];
380 let configuration: id = msg_send![configuration, init];
381
382 let frame: cocoa::foundation::NSRect = msg_send![sc_window, frame];
383 let width = frame.size.width as i64;
384 let height = frame.size.height as i64;
385 log::info!("capture_window_frame: window frame {}x{}", width, height);
386
387 if width <= 0 || height <= 0 {
388 if let Some(tx) = tx.borrow_mut().take() {
389 tx.send(Err(anyhow!(
390 "Window has invalid dimensions: {}x{}",
391 width,
392 height
393 )))
394 .ok();
395 }
396 return;
397 }
398
399 let _: () = msg_send![configuration, setWidth: width];
400 let _: () = msg_send![configuration, setHeight: height];
401 let _: () = msg_send![configuration, setScalesToFit: true];
402 let _: () = msg_send![configuration, setPixelFormat: 0x42475241u32]; // 'BGRA'
403 let _: () = msg_send![configuration, setShowsCursor: false];
404 let _: () = msg_send![configuration, setCapturesAudio: false];
405
406 let tx_for_capture = tx.clone();
407 // The completion handler receives (CGImageRef, NSError*), not CMSampleBuffer
408 let capture_handler =
409 ConcreteBlock::new(move |cg_image: core_graphics::sys::CGImageRef, error: id| {
410 log::info!("Screenshot capture handler called");
411
412 let Some(tx) = tx_for_capture.borrow_mut().take() else {
413 log::warn!("Screenshot capture: tx already taken");
414 return;
415 };
416
417 unsafe {
418 if error != nil {
419 let msg: id = msg_send![error, localizedDescription];
420 let error_str = NSStringExt::to_str(&msg);
421 log::error!("Screenshot capture error from API: {:?}", error_str);
422 tx.send(Err(anyhow!("Screenshot capture failed: {:?}", error_str)))
423 .ok();
424 return;
425 }
426
427 if cg_image.is_null() {
428 log::error!("Screenshot capture: cg_image is null");
429 tx.send(Err(anyhow!(
430 "Screenshot capture returned null CGImage. \
431 This may mean Screen Recording permission is not granted."
432 )))
433 .ok();
434 return;
435 }
436
437 log::info!("Screenshot capture: got CGImage, converting...");
438 let cg_image = CGImage::from_ptr(cg_image);
439 match cg_image_to_rgba_image(&cg_image) {
440 Ok(image) => {
441 log::info!(
442 "Screenshot capture: success! {}x{}",
443 image.width(),
444 image.height()
445 );
446 tx.send(Ok(image)).ok();
447 }
448 Err(e) => {
449 log::error!("Screenshot capture: CGImage conversion failed: {}", e);
450 tx.send(Err(e)).ok();
451 }
452 }
453 }
454 });
455 let capture_handler = capture_handler.copy();
456
457 log::info!("Calling SCScreenshotManager captureImageWithFilter...");
458 let _: () = msg_send![
459 class!(SCScreenshotManager),
460 captureImageWithFilter: filter
461 configuration: configuration
462 completionHandler: capture_handler
463 ];
464 log::info!("SCScreenshotManager captureImageWithFilter called");
465}
466
467/// Converts a CGImage to an RgbaImage.
468fn cg_image_to_rgba_image(cg_image: &CGImage) -> Result<RgbaImage> {
469 let width = cg_image.width();
470 let height = cg_image.height();
471
472 if width == 0 || height == 0 {
473 return Err(anyhow!("CGImage has zero dimensions: {}x{}", width, height));
474 }
475
476 // Create a bitmap context to draw the CGImage into
477 let color_space = CGColorSpace::create_device_rgb();
478 let bytes_per_row = width * 4;
479 let mut pixel_data: Vec<u8> = vec![0; height * bytes_per_row];
480
481 let context = core_graphics::context::CGContext::create_bitmap_context(
482 Some(pixel_data.as_mut_ptr() as *mut c_void),
483 width,
484 height,
485 8, // bits per component
486 bytes_per_row, // bytes per row
487 &color_space,
488 core_graphics::base::kCGImageAlphaPremultipliedLast // RGBA
489 | core_graphics::base::kCGBitmapByteOrder32Big,
490 );
491
492 // Draw the image into the context
493 let rect = core_graphics::geometry::CGRect::new(
494 &core_graphics::geometry::CGPoint::new(0.0, 0.0),
495 &core_graphics::geometry::CGSize::new(width as CGFloat, height as CGFloat),
496 );
497 context.draw_image(rect, cg_image);
498
499 // The pixel data is now in RGBA format
500 ImageBuffer::<Rgba<u8>, Vec<u8>>::from_raw(width as u32, height as u32, pixel_data)
501 .ok_or_else(|| anyhow!("Failed to create RgbaImage from CGImage pixel data"))
502}
503
504/// Converts a CVPixelBuffer (in BGRA format) to an RgbaImage.
505///
506/// This function locks the pixel buffer, reads the raw pixel data,
507/// converts from BGRA to RGBA format, and returns an image::RgbaImage.
508pub fn cv_pixel_buffer_to_rgba_image(pixel_buffer: &CVPixelBuffer) -> Result<RgbaImage> {
509 use core_video::r#return::kCVReturnSuccess;
510
511 unsafe {
512 if pixel_buffer.lock_base_address(0) != kCVReturnSuccess {
513 return Err(anyhow!("Failed to lock pixel buffer base address"));
514 }
515
516 let width = pixel_buffer.get_width();
517 let height = pixel_buffer.get_height();
518 let bytes_per_row = pixel_buffer.get_bytes_per_row();
519 let base_address = pixel_buffer.get_base_address();
520
521 if base_address.is_null() {
522 pixel_buffer.unlock_base_address(0);
523 return Err(anyhow!("Pixel buffer base address is null"));
524 }
525
526 let mut rgba_data = Vec::with_capacity(width * height * 4);
527
528 for y in 0..height {
529 let row_start = base_address.add(y * bytes_per_row) as *const u8;
530 for x in 0..width {
531 let pixel = row_start.add(x * 4);
532 let b = *pixel;
533 let g = *pixel.add(1);
534 let r = *pixel.add(2);
535 let a = *pixel.add(3);
536 rgba_data.push(r);
537 rgba_data.push(g);
538 rgba_data.push(b);
539 rgba_data.push(a);
540 }
541 }
542
543 pixel_buffer.unlock_base_address(0);
544
545 ImageBuffer::<Rgba<u8>, Vec<u8>>::from_raw(width as u32, height as u32, rgba_data)
546 .ok_or_else(|| anyhow!("Failed to create RgbaImage from pixel data"))
547 }
548}
549
550/// Converts a ScreenCaptureFrame to an RgbaImage.
551///
552/// This is useful for converting frames received from continuous screen capture streams.
553pub fn screen_capture_frame_to_rgba_image(frame: &ScreenCaptureFrame) -> Result<RgbaImage> {
554 unsafe {
555 let pixel_buffer =
556 CVPixelBuffer::wrap_under_get_rule(frame.0.as_concrete_TypeRef() as *mut _);
557 cv_pixel_buffer_to_rgba_image(&pixel_buffer)
558 }
559}
560
561#[ctor]
562unsafe fn build_classes() {
563 let mut decl = ClassDecl::new("GPUIStreamDelegate", class!(NSObject)).unwrap();
564 unsafe {
565 decl.add_method(
566 sel!(outputVideoEffectDidStartForStream:),
567 output_video_effect_did_start_for_stream as extern "C" fn(&Object, Sel, id),
568 );
569 decl.add_method(
570 sel!(outputVideoEffectDidStopForStream:),
571 output_video_effect_did_stop_for_stream as extern "C" fn(&Object, Sel, id),
572 );
573 decl.add_method(
574 sel!(stream:didStopWithError:),
575 stream_did_stop_with_error as extern "C" fn(&Object, Sel, id, id),
576 );
577 DELEGATE_CLASS = decl.register();
578
579 let mut decl = ClassDecl::new("GPUIStreamOutput", class!(NSObject)).unwrap();
580 decl.add_method(
581 sel!(stream:didOutputSampleBuffer:ofType:),
582 stream_did_output_sample_buffer_of_type
583 as extern "C" fn(&Object, Sel, id, id, NSInteger),
584 );
585 decl.add_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR);
586
587 OUTPUT_CLASS = decl.register();
588 }
589}
590
591extern "C" fn output_video_effect_did_start_for_stream(_this: &Object, _: Sel, _stream: id) {}
592
593extern "C" fn output_video_effect_did_stop_for_stream(_this: &Object, _: Sel, _stream: id) {}
594
595extern "C" fn stream_did_stop_with_error(_this: &Object, _: Sel, _stream: id, _error: id) {}
596
597extern "C" fn stream_did_output_sample_buffer_of_type(
598 this: &Object,
599 _: Sel,
600 _stream: id,
601 sample_buffer: id,
602 buffer_type: NSInteger,
603) {
604 if buffer_type != SCStreamOutputTypeScreen {
605 return;
606 }
607
608 unsafe {
609 let sample_buffer = sample_buffer as CMSampleBufferRef;
610 let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer);
611 if let Some(buffer) = sample_buffer.image_buffer() {
612 let callback: Box<Box<dyn Fn(ScreenCaptureFrame)>> =
613 Box::from_raw(*this.get_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR) as *mut _);
614 callback(ScreenCaptureFrame(buffer));
615 mem::forget(callback);
616 }
617 }
618}