From 1611635e5fd957809790f25dfd5a0426490d0aaf Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Tue, 30 Aug 2022 15:49:07 +0200 Subject: [PATCH] Capture display frames and access underlying IOSurface --- crates/capture/build.rs | 2 + crates/capture/src/main.rs | 156 ++++++++++++++++++++++++++++++++++++- 2 files changed, 155 insertions(+), 3 deletions(-) diff --git a/crates/capture/build.rs b/crates/capture/build.rs index 51ff36f705a1f969005c1bafc86fd560f446017f..550d1064ae7bcab0a8219ff1201e0071ccd847c1 100644 --- a/crates/capture/build.rs +++ b/crates/capture/build.rs @@ -24,6 +24,8 @@ fn main() { .clang_arg("-xobjective-c") .allowlist_function("CMTimeMake") .allowlist_type("SCStreamOutputType") + .allowlist_type("SCFrameStatus") + .allowlist_var("SCStreamFrameInfo.*") .allowlist_function("dispatch_queue_create") .parse_callbacks(Box::new(bindgen::CargoCallbacks)) .layout_tests(false) diff --git a/crates/capture/src/main.rs b/crates/capture/src/main.rs index c1c017c3599927f9eb6b823865e373b02491b44a..08a6cf03b2aaa9483b700aa96fde6e79a1ca413c 100644 --- a/crates/capture/src/main.rs +++ b/crates/capture/src/main.rs @@ -4,8 +4,10 @@ use crate::bindings::{dispatch_queue_create, NSObject, SCStreamOutputType}; use block::ConcreteBlock; use cocoa::{ base::{id, nil, YES}, - foundation::{NSArray, NSBundle, NSString, NSUInteger}, + foundation::{NSArray, NSString, NSUInteger}, }; +use core_foundation::{base::TCFType, number::CFNumberRef, string::CFStringRef}; +use core_media::{CMSampleBuffer, CMSampleBufferRef}; use gpui::{actions, elements::*, keymap::Binding, Menu, MenuItem}; use log::LevelFilter; use objc::{ @@ -49,6 +51,8 @@ fn main() { let applications: id = msg_send![content, applications]; let displays: id = msg_send![content, displays]; let display: id = displays.objectAtIndex(0); + let display_width: usize = msg_send![display, width]; + let display_height: usize = msg_send![display, height]; let mut decl = ClassDecl::new("CaptureOutput", class!(NSObject)).unwrap(); decl.add_protocol(Protocol::get("SCStreamOutput").unwrap()); @@ -63,6 +67,8 @@ fn main() { // let filter: id = msg_send![filter, initWithDesktopIndependentWindow: window]; let config: id = msg_send![class!(SCStreamConfiguration), alloc]; let config: id = msg_send![config, init]; + let _: () = msg_send![config, setWidth: display_width]; + let _: () = msg_send![config, setHeight: display_height]; let _: () = msg_send![config, setMinimumFrameInterval: bindings::CMTimeMake(1, 60)]; let _: () = msg_send![config, setQueueDepth: 6]; let _: () = msg_send![config, setShowsCursor: YES]; @@ -134,12 +140,156 @@ extern "C" fn sample_output( _: &Object, _: Sel, _stream: id, - _buffer: id, + buffer: id, _kind: SCStreamOutputType, ) { - println!("sample output"); + let buffer = unsafe { CMSampleBuffer::wrap_under_get_rule(buffer as CMSampleBufferRef) }; + let attachments = buffer.attachments(); + let attachments = attachments.first().expect("no attachments for sample"); + + unsafe { + let string = bindings::SCStreamFrameInfoStatus.0 as CFStringRef; + let status = core_foundation::number::CFNumber::wrap_under_get_rule( + *attachments.get(string) as CFNumberRef, + ) + .to_i64() + .expect("invalid frame info status"); + + if status != bindings::SCFrameStatus_SCFrameStatusComplete { + println!("received incomplete frame"); + return; + } + } + + let image_buffer = buffer.image_buffer(); + dbg!(image_buffer.width(), image_buffer.height()); + let io_surface = image_buffer.io_surface(); } fn quit(_: &Quit, cx: &mut gpui::MutableAppContext) { cx.platform().quit(); } + +mod core_media { + #![allow(non_snake_case)] + + use crate::core_video::{CVImageBuffer, CVImageBufferRef}; + use core_foundation::{ + array::{CFArray, CFArrayRef}, + base::{CFTypeID, TCFType}, + declare_TCFType, + dictionary::CFDictionary, + impl_CFTypeDescription, impl_TCFType, + string::CFString, + }; + use std::ffi::c_void; + + #[repr(C)] + pub struct __CMSampleBuffer(c_void); + // The ref type must be a pointer to the underlying struct. + pub type CMSampleBufferRef = *const __CMSampleBuffer; + + declare_TCFType!(CMSampleBuffer, CMSampleBufferRef); + impl_TCFType!(CMSampleBuffer, CMSampleBufferRef, CMSampleBufferGetTypeID); + impl_CFTypeDescription!(CMSampleBuffer); + + impl CMSampleBuffer { + pub fn attachments(&self) -> Vec> { + unsafe { + let attachments = + CMSampleBufferGetSampleAttachmentsArray(self.as_concrete_TypeRef(), true); + CFArray::::wrap_under_get_rule(attachments) + .into_iter() + .map(|attachments| { + CFDictionary::wrap_under_get_rule(attachments.as_concrete_TypeRef()) + }) + .collect() + } + } + + pub fn image_buffer(&self) -> CVImageBuffer { + unsafe { + CVImageBuffer::wrap_under_get_rule(CMSampleBufferGetImageBuffer( + self.as_concrete_TypeRef(), + )) + } + } + } + + extern "C" { + fn CMSampleBufferGetTypeID() -> CFTypeID; + fn CMSampleBufferGetSampleAttachmentsArray( + buffer: CMSampleBufferRef, + create_if_necessary: bool, + ) -> CFArrayRef; + fn CMSampleBufferGetImageBuffer(buffer: CMSampleBufferRef) -> CVImageBufferRef; + } +} + +mod core_video { + #![allow(non_snake_case)] + + use crate::io_surface::{IOSurface, IOSurfaceRef}; + use core_foundation::{ + base::{CFTypeID, TCFType}, + declare_TCFType, impl_CFTypeDescription, impl_TCFType, + }; + use std::ffi::c_void; + + #[repr(C)] + pub struct __CVImageBuffer(c_void); + // The ref type must be a pointer to the underlying struct. + pub type CVImageBufferRef = *const __CVImageBuffer; + + declare_TCFType!(CVImageBuffer, CVImageBufferRef); + impl_TCFType!(CVImageBuffer, CVImageBufferRef, CVImageBufferGetTypeID); + impl_CFTypeDescription!(CVImageBuffer); + + impl CVImageBuffer { + pub fn io_surface(&self) -> IOSurface { + unsafe { + IOSurface::wrap_under_get_rule(CVPixelBufferGetIOSurface( + self.as_concrete_TypeRef(), + )) + } + } + + pub fn width(&self) -> usize { + unsafe { CVPixelBufferGetWidth(self.as_concrete_TypeRef()) } + } + + pub fn height(&self) -> usize { + unsafe { CVPixelBufferGetHeight(self.as_concrete_TypeRef()) } + } + } + + extern "C" { + fn CVImageBufferGetTypeID() -> CFTypeID; + fn CVPixelBufferGetIOSurface(buffer: CVImageBufferRef) -> IOSurfaceRef; + fn CVPixelBufferGetWidth(buffer: CVImageBufferRef) -> usize; + fn CVPixelBufferGetHeight(buffer: CVImageBufferRef) -> usize; + } +} + +mod io_surface { + #![allow(non_snake_case)] + + use core_foundation::{ + base::{CFTypeID, TCFType}, + declare_TCFType, impl_CFTypeDescription, impl_TCFType, + }; + use std::ffi::c_void; + + #[repr(C)] + pub struct __IOSurface(c_void); + // The ref type must be a pointer to the underlying struct. + pub type IOSurfaceRef = *const __IOSurface; + + declare_TCFType!(IOSurface, IOSurfaceRef); + impl_TCFType!(IOSurface, IOSurfaceRef, IOSurfaceGetTypeID); + impl_CFTypeDescription!(IOSurface); + + extern "C" { + fn IOSurfaceGetTypeID() -> CFTypeID; + } +}