From b8e50ab47e005b6a7f3c8ebafdaac6fdd27a6b62 Mon Sep 17 00:00:00 2001
From: Andreas Reich <r_andreas2@web.de>
Date: Wed, 6 Nov 2024 01:46:00 +0100
Subject: [PATCH] Fix crash when using formats with bits per pixel size not
 divisible by 8 (#61)

* don't fail get_bytes_per_frame for byte aligned frame buffers
* handle failing to retrieve bytes more gracefully
---
 src/iter.rs    | 25 +++++++++++++++++++++----
 src/pix_fmt.rs |  7 +++++--
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/iter.rs b/src/iter.rs
index fe3ce5a..4f58dd3 100644
--- a/src/iter.rs
+++ b/src/iter.rs
@@ -214,10 +214,26 @@ pub fn spawn_stdout_thread(
     let mut buffers = stdout_output_video_streams
       .map(|video_stream| {
         // Since we filtered for video_streams above, we can unwrap unconditionally.
-        let video_data = video_stream.video_data().unwrap();
-        let bytes_per_frame = get_bytes_per_frame(video_data);
         let buf_size = match video_stream.format.as_str() {
-          "rawvideo" => bytes_per_frame.expect("Should use a known pix_fmt") as usize,
+          "rawvideo" => {
+            let Some(video_data) = video_stream.video_data() else {
+              tx.send(FfmpegEvent::Error(
+                "Video stream doesn't have any video data".to_owned(),
+              ))
+              .ok();
+              return Vec::new();
+            };
+
+            let Some(bytes_per_frame) = get_bytes_per_frame(video_data) else {
+              tx.send(FfmpegEvent::Error(
+                format!("Can't bytes per fraame for video data {video_data:?}").to_owned(),
+              ))
+              .ok();
+              return Vec::new();
+            };
+
+            bytes_per_frame as usize
+          }
 
           // Arbitrary default buffer size for receiving indeterminate chunks
           // of any encoder or container output, when frame boundaries are unknown
@@ -234,7 +250,8 @@ pub fn spawn_stdout_thread(
       })
       .collect::<Vec<Vec<u8>>>();
 
-    // No buffers probably indicates that output is being sent to file
+    // No buffers probably indicates that output is being sent to file or
+    // that an error occured.
     if buffers.is_empty() {
       return;
     }
diff --git a/src/pix_fmt.rs b/src/pix_fmt.rs
index 74c1b95..5577af9 100644
--- a/src/pix_fmt.rs
+++ b/src/pix_fmt.rs
@@ -234,8 +234,11 @@ pub fn get_bytes_per_frame(video_data: &VideoStream) -> Option<u32> {
   let bits_per_pixel = get_bits_per_pixel(&video_data.pix_fmt)?;
   // Enforce byte-alignment, since we don't currently have buffer reads in
   // sub-byte increments.
-  match bits_per_pixel % 8 {
-    0 => Some(video_data.width * video_data.height * bits_per_pixel / 8),
+  // Use the full frame buffer size for this, since formats like `yuvj420p` typically restrict a frame's size
+  // such that the buffer size has full bytes.
+  let num_bits = video_data.width * video_data.height * bits_per_pixel;
+  match num_bits % 8 {
+    0 => Some(num_bits / 8),
     _ => None,
   }
 }