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