From b51abc5a639883de5416f260cc146012068791d2 Mon Sep 17 00:00:00 2001
From: Nathan Sobo <nathan@zed.dev>
Date: Wed, 31 Aug 2022 13:36:29 -0600
Subject: [PATCH] Read the frame data out of the CMSampleBuffer

Still not sending it anywhere, but think I'm reading it correctly.
---
 Cargo.lock                    |  1 +
 crates/capture/Cargo.toml     |  5 +---
 crates/capture/script/capture |  5 ----
 crates/capture/src/main.rs    | 20 +++++++++++++++
 crates/media/src/media.rs     | 48 ++++++++++++++++++++++++++++++++++-
 5 files changed, 69 insertions(+), 10 deletions(-)
 delete mode 100755 crates/capture/script/capture

diff --git a/Cargo.lock b/Cargo.lock
index 6b34c02f0d..60e54aeb9f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -757,6 +757,7 @@ dependencies = [
  "anyhow",
  "bindgen",
  "block",
+ "byteorder",
  "bytes",
  "cocoa",
  "core-foundation",
diff --git a/crates/capture/Cargo.toml b/crates/capture/Cargo.toml
index c2bee8312a..cdae75fc2b 100644
--- a/crates/capture/Cargo.toml
+++ b/crates/capture/Cargo.toml
@@ -4,10 +4,6 @@ version = "0.1.0"
 edition = "2021"
 description = "An example of screen capture"
 
-[package.metadata.bundle]
-name = "Capture"
-identifier = "dev.zed.Capture"
-
 [dependencies]
 gpui = { path = "../gpui" }
 media = { path = "../media" }
@@ -15,6 +11,7 @@ media = { path = "../media" }
 anyhow = "1.0.38"
 block = "0.1"
 bytes = "1.2"
+byteorder = "1.4"
 cocoa = "0.24"
 core-foundation = "0.9.3"
 core-graphics = "0.22.3"
diff --git a/crates/capture/script/capture b/crates/capture/script/capture
deleted file mode 100755
index 0e6cf1f409..0000000000
--- a/crates/capture/script/capture
+++ /dev/null
@@ -1,5 +0,0 @@
-#!/bin/bash
-
-cargo bundle
-TTY=`tty`
-open ../../target/debug/bundle/osx/Capture.app --stdout $TTY --stderr $TTY
diff --git a/crates/capture/src/main.rs b/crates/capture/src/main.rs
index b1d8f11795..97a5557c66 100644
--- a/crates/capture/src/main.rs
+++ b/crates/capture/src/main.rs
@@ -3,6 +3,7 @@ mod compression_session;
 
 use crate::{bindings::SCStreamOutputType, compression_session::CompressionSession};
 use block::ConcreteBlock;
+use byteorder::{BigEndian, ReadBytesExt};
 use bytes::BytesMut;
 use cocoa::{
     base::{id, nil, YES},
@@ -128,6 +129,25 @@ impl ScreenCaptureView {
                                 let nal_unit = compression_buffer.split();
                             }
                         }
+
+                        let data = sample_buffer.data();
+                        let mut data = data.bytes();
+
+                        const AVCC_HEADER_LENGTH: usize = 4;
+                        while data.len() - AVCC_HEADER_LENGTH > 0 {
+                            let nal_unit_len = match data.read_u32::<BigEndian>() {
+                                Ok(len) => len as usize,
+                                Err(error) => {
+                                    log::error!("error decoding nal unit length: {}", error);
+                                    return;
+                                }
+                            };
+                            compression_buffer.extend_from_slice(&START_CODE);
+                            compression_buffer.extend_from_slice(&data[..nal_unit_len as usize]);
+                            data = &data[nal_unit_len..];
+
+                            let nal_unit = compression_buffer.split();
+                        }
                     },
                 )
                 .unwrap();
diff --git a/crates/media/src/media.rs b/crates/media/src/media.rs
index 5c5f43af9d..ebe3ef7f4d 100644
--- a/crates/media/src/media.rs
+++ b/crates/media/src/media.rs
@@ -269,6 +269,14 @@ pub mod core_media {
                 ))
             }
         }
+
+        pub fn data(&self) -> CMBlockBuffer {
+            unsafe {
+                CMBlockBuffer::wrap_under_get_rule(CMSampleBufferGetDataBuffer(
+                    self.as_concrete_TypeRef(),
+                ))
+            }
+        }
     }
 
     #[link(name = "CoreMedia", kind = "framework")]
@@ -285,11 +293,11 @@ pub mod core_media {
             timing_info_out: *mut CMSampleTimingInfo,
         ) -> OSStatus;
         fn CMSampleBufferGetFormatDescription(buffer: CMSampleBufferRef) -> CMFormatDescriptionRef;
+        fn CMSampleBufferGetDataBuffer(sample_buffer: CMSampleBufferRef) -> CMBlockBufferRef;
     }
 
     #[repr(C)]
     pub struct __CMFormatDescription(c_void);
-    // The ref type must be a pointer to the underlying struct.
     pub type CMFormatDescriptionRef = *const __CMFormatDescription;
 
     declare_TCFType!(CMFormatDescription, CMFormatDescriptionRef);
@@ -350,6 +358,44 @@ pub mod core_media {
             NALUnitHeaderLengthOut: *mut isize,
         ) -> OSStatus;
     }
+
+    #[repr(C)]
+    pub struct __CMBlockBuffer(c_void);
+    pub type CMBlockBufferRef = *const __CMBlockBuffer;
+
+    declare_TCFType!(CMBlockBuffer, CMBlockBufferRef);
+    impl_TCFType!(CMBlockBuffer, CMBlockBufferRef, CMBlockBufferGetTypeID);
+    impl_CFTypeDescription!(CMBlockBuffer);
+
+    impl CMBlockBuffer {
+        pub fn bytes(&self) -> &[u8] {
+            unsafe {
+                let mut bytes = ptr::null();
+                let mut len = 0;
+                let result = CMBlockBufferGetDataPointer(
+                    self.as_concrete_TypeRef(),
+                    0,
+                    &mut 0,
+                    &mut len,
+                    &mut bytes,
+                );
+                assert!(result == 0, "could not get block buffer data");
+                std::slice::from_raw_parts(bytes, len)
+            }
+        }
+    }
+
+    #[link(name = "CoreMedia", kind = "framework")]
+    extern "C" {
+        fn CMBlockBufferGetTypeID() -> CFTypeID;
+        fn CMBlockBufferGetDataPointer(
+            buffer: CMBlockBufferRef,
+            offset: usize,
+            length_at_offset_out: *mut usize,
+            total_length_out: *mut usize,
+            data_pointer_out: *mut *const u8,
+        ) -> OSStatus;
+    }
 }
 
 pub mod video_toolbox {