use heap allocated xcoder params to prevent stack overflow

sportsball-ai · Jul 15, 2024 · bfae5c3 · bfae5c3
1 parent 6f105f3
commit bfae5c3
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 19 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/mpegts-segmenter/src/analyzer.rs b/mpegts-segmenter/src/analyzer.rs
@@ -2,10 +2,7 @@ use mpeg2::{pes, ts};
 use mpeg4::AudioDataTransportStream;
 use vecmap::VecMap;
 
-use std::{
-    collections::VecDeque,
-    error::Error,
-};
+use std::{collections::VecDeque, error::Error};
 
 pub type Result<T> = std::result::Result<T, Box<dyn Error + Send + Sync>>;
 

diff --git a/xcoder/xcoder-quadra/Cargo.toml b/xcoder/xcoder-quadra/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2021"
 av-traits = { path = "../../av-traits" }
 snafu = "0.8.0"
 scopeguard = "1.1.0"
+libc = "0.2"
 
 [dev-dependencies]
 h264 = { path = "../../h264" }

diff --git a/xcoder/xcoder-quadra/src/decoder.rs b/xcoder/xcoder-quadra/src/decoder.rs
@@ -1,7 +1,12 @@
-use super::{fps_to_rational, XcoderHardware, XcoderHardwareFrame};
+use super::{alloc_zeroed, fps_to_rational, XcoderHardware, XcoderHardwareFrame};
 use scopeguard::{guard, ScopeGuard};
 use snafu::{Error, Snafu};
-use std::{marker::PhantomData, mem, ops::Deref};
+use std::{
+    marker::PhantomData,
+    mem::{self, MaybeUninit},
+    ops::Deref,
+    os::raw::c_void,
+};
 use xcoder_quadra_sys as sys;
 
 #[derive(Clone, Copy, Debug)]
@@ -205,17 +210,17 @@ impl<E: Error, I: XcoderDecoderInput<E>> XcoderDecoder<I, E> {
         let (fps_numerator, fps_denominator) = fps_to_rational(config.fps);
 
         unsafe {
-            let mut params: sys::ni_xcoder_params_t = mem::zeroed();
-            let code = sys::ni_decoder_init_default_params(&mut params as _, fps_numerator, fps_denominator, 0, config.width, config.height);
+            let mut params = alloc_zeroed();
+            let code = sys::ni_decoder_init_default_params(params.as_mut_ptr(), fps_numerator, fps_denominator, 0, config.width, config.height);
             if code != sys::ni_retcode_t_NI_RETCODE_SUCCESS {
                 return Err(XcoderDecoderError::Unknown {
                     code,
                     operation: "initializing parameters",
                 });
             }
+            let mut params = mem::transmute::<Box<MaybeUninit<sys::ni_xcoder_params_t>>, Box<sys::ni_xcoder_params_t>>(params);
             params.__bindgen_anon_1.dec_input_params.hwframes = 1;
             params.__bindgen_anon_1.dec_input_params.mcmode = config.multicore_joint_mode.into();
-            let params = Box::new(params);
 
             let session = sys::ni_device_session_context_alloc_init();
             if session.is_null() {
@@ -227,7 +232,7 @@ impl<E: Error, I: XcoderDecoderInput<E>> XcoderDecoder<I, E> {
 
             (**session).hw_id = config.hardware_id.unwrap_or(-1);
             (**session).hw_action = sys::ni_codec_hw_actions_NI_CODEC_HW_ENABLE as _;
-            (**session).p_session_config = &*params as *const sys::ni_xcoder_params_t as _;
+            (**session).p_session_config = params.as_mut() as *mut sys::ni_xcoder_params_t as *mut c_void;
             (**session).codec_format = match config.codec {
                 XcoderDecoderCodec::H264 => sys::_ni_codec_format_NI_CODEC_FORMAT_H264,
                 XcoderDecoderCodec::H265 => sys::_ni_codec_format_NI_CODEC_FORMAT_H265,

diff --git a/xcoder/xcoder-quadra/src/encoder.rs b/xcoder/xcoder-quadra/src/encoder.rs
@@ -1,8 +1,12 @@
-use super::{fps_to_rational, XcoderHardware, XcoderHardwareFrame};
+use super::{alloc_zeroed, fps_to_rational, XcoderHardware, XcoderHardwareFrame};
 use av_traits::{EncodedFrameType, EncodedVideoFrame, RawVideoFrame, VideoEncoder, VideoEncoderOutput};
 use scopeguard::{guard, ScopeGuard};
 use snafu::Snafu;
-use std::{collections::VecDeque, mem};
+use std::{
+    collections::VecDeque,
+    mem::{self, MaybeUninit},
+    os::raw::c_void,
+};
 use xcoder_quadra_sys as sys;
 
 #[derive(Debug, Snafu)]
@@ -46,7 +50,7 @@ impl EncodedFrame {
 
     pub fn as_slice(&self) -> &[u8] {
         let data = unsafe { &std::slice::from_raw_parts(self.data_io.data.packet.p_data as _, self.data_io.data.packet.data_len as _) };
-        &data[self.meta_size as usize..]
+        &data[self.meta_size..]
     }
 
     /// When parameter sets are emitted by the encoder, they're provided here.
@@ -141,9 +145,9 @@ impl<F> XcoderEncoder<F> {
         let (fps_numerator, fps_denominator) = fps_to_rational(config.fps);
 
         unsafe {
-            let mut params: sys::ni_xcoder_params_t = mem::zeroed();
+            let mut params = alloc_zeroed::<sys::ni_xcoder_params_t>();
             let code = sys::ni_encoder_init_default_params(
-                &mut params as _,
+                params.as_mut_ptr(),
                 fps_numerator,
                 fps_denominator,
                 // There's nothing special about this default bitrate. It's not used unless
@@ -162,7 +166,7 @@ impl<F> XcoderEncoder<F> {
                     operation: "initializing parameters",
                 });
             }
-
+            let mut params = mem::transmute::<Box<MaybeUninit<sys::ni_xcoder_params_t>>, Box<sys::ni_xcoder_params_t>>(params);
             let cfg_enc_params = &mut params.__bindgen_anon_1.cfg_enc_params;
             cfg_enc_params.planar = 1;
             cfg_enc_params.rc.enable_rate_control = if config.bitrate.is_some() { 1 } else { 0 };
@@ -234,8 +238,6 @@ impl<F> XcoderEncoder<F> {
                 params.hwframes = 1;
             }
 
-            let params = Box::new(params);
-
             let session = sys::ni_device_session_context_alloc_init();
             if session.is_null() {
                 return Err(XcoderEncoderError::UnableToAllocateDeviceSessionContext);
@@ -250,7 +252,7 @@ impl<F> XcoderEncoder<F> {
                 (**session).sender_handle = hw.device_handle;
                 (**session).hw_action = sys::ni_codec_hw_actions_NI_CODEC_HW_ENABLE as _;
             }
-            (**session).p_session_config = &*params as *const sys::ni_xcoder_params_t as _;
+            (**session).p_session_config = params.as_mut() as *mut sys::ni_xcoder_params_t as *mut c_void;
             (**session).codec_format = match config.codec {
                 XcoderEncoderCodec::H264 { .. } => sys::_ni_codec_format_NI_CODEC_FORMAT_H264,
                 XcoderEncoderCodec::H265 { .. } => sys::_ni_codec_format_NI_CODEC_FORMAT_H265,

diff --git a/xcoder/xcoder-quadra/src/lib.rs b/xcoder/xcoder-quadra/src/lib.rs
@@ -90,6 +90,17 @@ mod linux_impl {
         let num = (fps * den as f64).round() as _;
         (num, den)
     }
+
+    fn alloc_zeroed<T>() -> Box<std::mem::MaybeUninit<T>> {
+        use std::alloc::GlobalAlloc as _;
+        unsafe {
+            let raw = std::alloc::System.alloc_zeroed(std::alloc::Layout::from_size_align_unchecked(
+                std::mem::size_of::<T>(),
+                std::mem::align_of::<T>(),
+            ));
+            Box::from_raw(raw as *mut std::mem::MaybeUninit<T>)
+        }
+    }
 }
 
 #[cfg(target_os = "linux")]