From aa1ce3edcad0fcd286d39b3e0c2fdc610c3988e7 Mon Sep 17 00:00:00 2001 From: Kieran Date: Wed, 11 Jun 2025 23:57:40 +0100 Subject: [PATCH] fix: decoder state --- src/decode.rs | 91 +++++++++++++++++++++++++++++++++++++++++---------- src/demux.rs | 28 +++++++++------- src/lib.rs | 13 ++++++-- 3 files changed, 100 insertions(+), 32 deletions(-) diff --git a/src/decode.rs b/src/decode.rs index 2f891c7..4664b9a 100644 --- a/src/decode.rs +++ b/src/decode.rs @@ -1,10 +1,10 @@ -use crate::{bail_ffmpeg, options_to_dict, rstr, StreamInfo}; +use crate::{bail_ffmpeg, rstr, StreamInfo}; use std::collections::hash_map::Entry; use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; use std::ptr; -use anyhow::Error; +use anyhow::{bail, Error}; use ffmpeg_sys_the_third::{ av_buffer_ref, av_frame_alloc, av_frame_free, av_hwdevice_ctx_create, av_hwdevice_get_type_name, av_hwdevice_iterate_types, avcodec_alloc_context3, @@ -50,6 +50,8 @@ impl Drop for DecoderCodecContext { if !self.context.is_null() { avcodec_free_context(&mut self.context); } + self.context = ptr::null_mut(); + self.codec = ptr::null_mut(); } } } @@ -166,30 +168,60 @@ impl Decoder { "Codec parameters are missing from stream" ); - let ctx = self.setup_decoder_codec((*codec_par).codec_id, (*stream).index, options)?; + let ctx = self.add_decoder((*codec_par).codec_id, (*stream).index)?; let ret = avcodec_parameters_to_context(ctx.context, (*stream).codecpar); bail_ffmpeg!(ret, "Failed to copy codec parameters to context"); - Ok(ctx) + + let stream_index = (*stream).index; + self.open_decoder_codec_by_index(stream_index, options)?; + Ok(self.codecs.get_mut(&stream_index).unwrap()) + } + + /// Open a decoder codec after parameters are set + pub unsafe fn open_decoder_codec(&mut self, ctx: &DecoderCodecContext) -> Result<(), Error> { + let mut dict = ptr::null_mut(); + let ret = avcodec_open2(ctx.context, ctx.codec, &mut dict); + bail_ffmpeg!(ret, "Failed to open codec"); + Ok(()) + } + + /// Open a decoder codec by stream index + pub unsafe fn open_decoder_codec_by_index( + &mut self, + stream_index: i32, + options: Option>, + ) -> Result<(), Error> { + if let Some(ctx) = self.codecs.get(&stream_index) { + let mut dict = if let Some(options) = options { + crate::options_to_dict(options)? + } else { + ptr::null_mut() + }; + let ret = avcodec_open2(ctx.context, ctx.codec, &mut dict); + bail_ffmpeg!(ret, "Failed to open codec"); + Ok(()) + } else { + bail!("Decoder not found for stream index {}", stream_index) + } } /// Configure a decoder manually - pub unsafe fn setup_decoder_codec( + pub unsafe fn add_decoder( &mut self, codec_id: AVCodecID, stream_index: i32, - options: Option>, ) -> Result<&mut DecoderCodecContext, Error> { if let Entry::Vacant(e) = self.codecs.entry(stream_index) { let codec = avcodec_find_decoder(codec_id); if codec.is_null() { - anyhow::bail!( + bail!( "Failed to find codec: {}", rstr!(avcodec_get_name(codec_id)) ) } let context = avcodec_alloc_context3(codec); if context.is_null() { - anyhow::bail!("Failed to alloc context") + bail!("Failed to alloc context") } let mut ret = 0; @@ -225,15 +257,6 @@ impl Decoder { } } } - let mut dict = if let Some(options) = options { - options_to_dict(options)? - } else { - ptr::null_mut() - }; - - ret = avcodec_open2(context, codec, &mut dict); - bail_ffmpeg!(ret, "Failed to open codec"); - let ctx = DecoderCodecContext { context, codec, @@ -243,7 +266,7 @@ impl Decoder { trace!("setup decoder={}", ctx); Ok(e.insert(ctx)) } else { - anyhow::bail!("Decoder already setup"); + bail!("Decoder already setup"); } } @@ -293,3 +316,35 @@ impl Decoder { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::Demuxer; + use ffmpeg_sys_the_third::av_packet_free; + + #[test] + fn decode_files() -> anyhow::Result<()> { + unsafe { + let files = std::fs::read_dir("./test_output/").unwrap(); + for file in files.into_iter() { + let mut mux = Demuxer::new(file.unwrap().path().to_str().unwrap())?; + let probe = mux.probe_input()?; + let mut decoder = Decoder::new(); + for stream in probe.streams.iter() { + decoder.setup_decoder(stream, None)?; + } + loop { + let (mut pkt, _) = mux.get_packet()?; + if pkt.is_null() { + break; + } + + decoder.decode_pkt(pkt)?; + av_packet_free(&mut pkt); + } + } + } + Ok(()) + } +} diff --git a/src/demux.rs b/src/demux.rs index 398efde..fda64f1 100644 --- a/src/demux.rs +++ b/src/demux.rs @@ -51,7 +51,7 @@ impl Demuxer { Ok(Self { ctx, input: DemuxerInput::Url(input.to_string()), - buffer_size: 1024 * 16, + buffer_size: 4096, }) } } @@ -92,12 +92,14 @@ impl Demuxer { unsafe fn open(&mut self) -> Result<()> { match &mut self.input { DemuxerInput::Url(input) => { + let input_cstr = cstr!(input.as_str()); let ret = avformat_open_input( &mut self.ctx, - cstr!(input.as_str()), + input_cstr, ptr::null_mut(), ptr::null_mut(), ); + libc::free(input_cstr as *mut libc::c_void); bail_ffmpeg!(ret); Ok(()) } @@ -117,16 +119,20 @@ impl Demuxer { } (*self.ctx).pb = pb; + let url_cstr = if let Some(url) = url { + cstr!(url.as_str()) + } else { + ptr::null_mut() + }; let ret = avformat_open_input( &mut self.ctx, - if let Some(url) = url { - cstr!(url.as_str()) as _ - } else { - ptr::null_mut() - }, + url_cstr, ptr::null_mut(), ptr::null_mut(), ); + if !url_cstr.is_null() { + libc::free(url_cstr as *mut libc::c_void); + } bail_ffmpeg!(ret); Ok(()) } @@ -176,7 +182,9 @@ impl Demuxer { while n_stream < (*self.ctx).nb_streams as usize { let stream = *(*self.ctx).streams.add(n_stream); n_stream += 1; - let lang = av_dict_get((*stream).metadata, cstr!("language"), ptr::null_mut(), 0); + let lang_key = cstr!("language"); + let lang = av_dict_get((*stream).metadata, lang_key, ptr::null_mut(), 0); + libc::free(lang_key as *mut libc::c_void); let language = if lang.is_null() { "".to_string() } else { @@ -261,8 +269,7 @@ impl Demuxer { let stream = self.get_stream((*pkt).stream_index as _)?; (*pkt).time_base = (*stream).time_base; - let pkg = (pkt, stream); - Ok(pkg) + Ok((pkt, stream)) } /// Get stream by index from context @@ -301,7 +308,6 @@ mod tests { #[cfg(feature = "avformat_version_greater_than_60_19")] #[test] - #[ignore] fn test_stream_groups() -> Result<()> { unsafe { let mut demux = diff --git a/src/lib.rs b/src/lib.rs index 6d0689e..0865ac6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -52,7 +52,6 @@ macro_rules! bail_ffmpeg { #[macro_export] macro_rules! cstr { ($str:expr) => { - // TODO: leaky std::ffi::CString::new($str).unwrap().into_raw() }; } @@ -132,7 +131,11 @@ fn get_ffmpeg_error_msg(ret: libc::c_int) -> String { unsafe fn options_to_dict(options: HashMap) -> Result<*mut AVDictionary, Error> { let mut dict = ptr::null_mut(); for (key, value) in options { - let ret = av_dict_set(&mut dict, cstr!(key), cstr!(value), 0); + let key_cstr = cstr!(key); + let value_cstr = cstr!(value); + let ret = av_dict_set(&mut dict, key_cstr, value_cstr, 0); + libc::free(key_cstr as *mut libc::c_void); + libc::free(value_cstr as *mut libc::c_void); bail_ffmpeg!(ret); } Ok(dict) @@ -181,7 +184,11 @@ fn list_opts(ctx: *mut libc::c_void) -> Result, Error> { fn set_opts(ctx: *mut libc::c_void, options: HashMap) -> Result<(), Error> { unsafe { for (key, value) in options { - let ret = av_opt_set(ctx, cstr!(key), cstr!(value), AV_OPT_SEARCH_CHILDREN); + let key_cstr = cstr!(key); + let value_cstr = cstr!(value); + let ret = av_opt_set(ctx, key_cstr, value_cstr, AV_OPT_SEARCH_CHILDREN); + libc::free(key_cstr as *mut libc::c_void); + libc::free(value_cstr as *mut libc::c_void); bail_ffmpeg!(ret); } }