scuffle_ffmpeg/io/
output.rs

1use std::ffi::CString;
2use std::ptr::NonNull;
3
4use super::internal::{Inner, InnerOptions, seek, write_packet};
5use crate::consts::DEFAULT_BUFFER_SIZE;
6use crate::dict::Dictionary;
7use crate::error::{FfmpegError, FfmpegErrorCode};
8use crate::ffi::*;
9use crate::packet::Packet;
10use crate::stream::Stream;
11use crate::{AVFmtFlags, AVFormatFlags};
12
13/// A struct that represents the options for the output.
14#[derive(Debug, Clone, bon::Builder)]
15pub struct OutputOptions {
16    /// The buffer size for the output.
17    #[builder(default = DEFAULT_BUFFER_SIZE)]
18    buffer_size: usize,
19    #[builder(setters(vis = "", name = format_ffi_internal))]
20    format_ffi: *const AVOutputFormat,
21}
22
23impl<S: output_options_builder::State> OutputOptionsBuilder<S> {
24    /// Sets the format FFI.
25    ///
26    /// Returns an error if the format FFI is null.
27    pub fn format_ffi(
28        self,
29        format_ffi: *const AVOutputFormat,
30    ) -> Result<OutputOptionsBuilder<output_options_builder::SetFormatFfi<S>>, FfmpegError>
31    where
32        S::FormatFfi: output_options_builder::IsUnset,
33    {
34        if format_ffi.is_null() {
35            return Err(FfmpegError::Arguments("could not determine output format"));
36        }
37
38        Ok(self.format_ffi_internal(format_ffi))
39    }
40
41    /// Gets the format ffi from the format name.
42    ///
43    /// Returns an error if the format name is empty or the format was not found.
44    #[inline]
45    pub fn format_name(
46        self,
47        format_name: &str,
48    ) -> Result<OutputOptionsBuilder<output_options_builder::SetFormatFfi<S>>, FfmpegError>
49    where
50        S::FormatFfi: output_options_builder::IsUnset,
51    {
52        self.format_name_mime_type(format_name, "")
53    }
54
55    /// Gets the format ffi from the format mime type.
56    ///
57    /// Returns an error if the format mime type is empty or the format was not found.
58    #[inline]
59    pub fn format_mime_type(
60        self,
61        format_mime_type: &str,
62    ) -> Result<OutputOptionsBuilder<output_options_builder::SetFormatFfi<S>>, FfmpegError>
63    where
64        S::FormatFfi: output_options_builder::IsUnset,
65    {
66        self.format_name_mime_type("", format_mime_type)
67    }
68
69    /// Sets the format ffi from the format name and mime type.
70    ///
71    /// Returns an error if both the format name and mime type are empty or the format was not found.
72    pub fn format_name_mime_type(
73        self,
74        format_name: &str,
75        format_mime_type: &str,
76    ) -> Result<OutputOptionsBuilder<output_options_builder::SetFormatFfi<S>>, FfmpegError>
77    where
78        S::FormatFfi: output_options_builder::IsUnset,
79    {
80        let c_format_name = CString::new(format_name).ok();
81        let c_format_mime_type = CString::new(format_mime_type).ok();
82        let c_format_name_ptr = c_format_name.as_ref().map(|s| s.as_ptr()).unwrap_or(std::ptr::null());
83        let c_format_mime_type_ptr = c_format_mime_type.as_ref().map(|s| s.as_ptr()).unwrap_or(std::ptr::null());
84        // Safety: av_guess_format is safe to call and all the arguments are valid
85        let format_ffi = unsafe { av_guess_format(c_format_name_ptr, std::ptr::null(), c_format_mime_type_ptr) };
86        self.format_ffi(format_ffi)
87    }
88}
89
90/// A struct that represents the output.
91pub struct Output<T: Send + Sync> {
92    inner: Inner<T>,
93    state: OutputState,
94}
95
96/// Safety: `T` must be `Send` and `Sync`.
97unsafe impl<T: Send + Sync> Send for Output<T> {}
98
99#[derive(Debug, Clone, Copy, PartialEq, Eq)]
100enum OutputState {
101    Uninitialized,
102    HeaderWritten,
103    TrailerWritten,
104}
105
106impl<T: Send + Sync> Output<T> {
107    /// Consumes the `Output` and returns the inner data.
108    pub fn into_inner(mut self) -> T {
109        *(self.inner.data.take().unwrap())
110    }
111}
112
113impl<T: std::io::Write + Send + Sync> Output<T> {
114    /// Creates a new `Output` with the given output and options.
115    pub fn new(output: T, options: OutputOptions) -> Result<Self, FfmpegError> {
116        Ok(Self {
117            inner: Inner::new(
118                output,
119                InnerOptions {
120                    buffer_size: options.buffer_size,
121                    write_fn: Some(write_packet::<T>),
122                    output_format: options.format_ffi,
123                    ..Default::default()
124                },
125            )?,
126            state: OutputState::Uninitialized,
127        })
128    }
129
130    /// Creates a new `Output` with the given output and options. The output must be seekable.
131    pub fn seekable(output: T, options: OutputOptions) -> Result<Self, FfmpegError>
132    where
133        T: std::io::Seek,
134    {
135        Ok(Self {
136            inner: Inner::new(
137                output,
138                InnerOptions {
139                    buffer_size: options.buffer_size,
140                    write_fn: Some(write_packet::<T>),
141                    seek_fn: Some(seek::<T>),
142                    output_format: options.format_ffi,
143                    ..Default::default()
144                },
145            )?,
146            state: OutputState::Uninitialized,
147        })
148    }
149}
150
151impl<T: Send + Sync> Output<T> {
152    /// Sets the metadata for the output.
153    pub fn set_metadata(&mut self, metadata: Dictionary) {
154        // Safety: We want to replace the metadata from the context (if one exists). This is safe as the metadata should be a valid pointer.
155        unsafe {
156            Dictionary::from_ptr_owned(self.inner.context.as_deref_mut_except().metadata);
157        };
158
159        self.inner.context.as_deref_mut_except().metadata = metadata.leak();
160    }
161
162    /// Returns the pointer to the underlying AVFormatContext.
163    pub const fn as_ptr(&self) -> *const AVFormatContext {
164        self.inner.context.as_ptr()
165    }
166
167    /// Returns the pointer to the underlying AVFormatContext.
168    pub const fn as_mut_ptr(&mut self) -> *mut AVFormatContext {
169        self.inner.context.as_mut_ptr()
170    }
171
172    /// Adds a new stream to the output.
173    pub fn add_stream(&mut self, codec: Option<*const AVCodec>) -> Option<Stream<'_>> {
174        let mut stream =
175            // Safety: `avformat_new_stream` is safe to call.
176            NonNull::new(unsafe { avformat_new_stream(self.as_mut_ptr(), codec.unwrap_or_else(std::ptr::null)) })?;
177
178        // Safety: The stream is a valid non-null pointer.
179        let stream = unsafe { stream.as_mut() };
180        stream.id = self.inner.context.as_deref_except().nb_streams as i32 - 1;
181
182        Some(Stream::new(stream, self.inner.context.as_mut_ptr()))
183    }
184
185    /// Copies a stream from the input to the output.
186    pub fn copy_stream<'a>(&'a mut self, stream: &Stream<'_>) -> Result<Option<Stream<'a>>, FfmpegError> {
187        let Some(codec_param) = stream.codec_parameters() else {
188            return Ok(None);
189        };
190
191        // Safety: `avformat_new_stream` is safe to call.
192        let Some(mut out_stream) = NonNull::new(unsafe { avformat_new_stream(self.as_mut_ptr(), std::ptr::null()) }) else {
193            return Ok(None);
194        };
195
196        // Safety: The stream is a valid non-null pointer.
197        let out_stream = unsafe { out_stream.as_mut() };
198
199        // Safety: `avcodec_parameters_copy` is safe to call when all arguments are valid.
200        FfmpegErrorCode(unsafe { avcodec_parameters_copy(out_stream.codecpar, codec_param) }).result()?;
201
202        out_stream.id = self.inner.context.as_deref_except().nb_streams as i32 - 1;
203
204        let mut out_stream = Stream::new(out_stream, self.inner.context.as_mut_ptr());
205        out_stream.set_time_base(stream.time_base());
206        out_stream.set_start_time(stream.start_time());
207        out_stream.set_duration(stream.duration());
208
209        Ok(Some(out_stream))
210    }
211
212    /// Writes the header to the output.
213    pub fn write_header(&mut self) -> Result<(), FfmpegError> {
214        if self.state != OutputState::Uninitialized {
215            return Err(FfmpegError::Arguments("header already written"));
216        }
217
218        // Safety: `avformat_write_header` is safe to call, if the header has not been
219        // written yet.
220        FfmpegErrorCode(unsafe { avformat_write_header(self.as_mut_ptr(), std::ptr::null_mut()) }).result()?;
221        self.state = OutputState::HeaderWritten;
222
223        Ok(())
224    }
225
226    /// Writes the header to the output with the given options.
227    pub fn write_header_with_options(&mut self, options: &mut Dictionary) -> Result<(), FfmpegError> {
228        if self.state != OutputState::Uninitialized {
229            return Err(FfmpegError::Arguments("header already written"));
230        }
231
232        // Safety: `avformat_write_header` is safe to call, if the header has not been
233        // written yet.
234        FfmpegErrorCode(unsafe { avformat_write_header(self.as_mut_ptr(), options.as_mut_ptr_ref()) }).result()?;
235        self.state = OutputState::HeaderWritten;
236
237        Ok(())
238    }
239
240    /// Writes the trailer to the output.
241    pub fn write_trailer(&mut self) -> Result<(), FfmpegError> {
242        if self.state != OutputState::HeaderWritten {
243            return Err(FfmpegError::Arguments(
244                "cannot write trailer before header or after trailer has been written",
245            ));
246        }
247
248        // Safety: `av_write_trailer` is safe to call, once the header has been written.
249        FfmpegErrorCode(unsafe { av_write_trailer(self.as_mut_ptr()) }).result()?;
250        self.state = OutputState::TrailerWritten;
251
252        Ok(())
253    }
254
255    /// Writes the interleaved packet to the output.
256    /// The difference between this and `write_packet` is that this function
257    /// writes the packet to the output and reorders the packets based on the
258    /// dts and pts.
259    pub fn write_interleaved_packet(&mut self, mut packet: Packet) -> Result<(), FfmpegError> {
260        if self.state != OutputState::HeaderWritten {
261            return Err(FfmpegError::Arguments(
262                "cannot write interleaved packet before header or after trailer has been written",
263            ));
264        }
265
266        // Safety: `av_interleaved_write_frame` is safe to call, once the header has
267        // been written.
268        FfmpegErrorCode(unsafe { av_interleaved_write_frame(self.as_mut_ptr(), packet.as_mut_ptr()) }).result()?;
269        Ok(())
270    }
271
272    /// Writes the packet to the output. Without reordering the packets.
273    pub fn write_packet(&mut self, packet: &Packet) -> Result<(), FfmpegError> {
274        if self.state != OutputState::HeaderWritten {
275            return Err(FfmpegError::Arguments(
276                "cannot write packet before header or after trailer has been written",
277            ));
278        }
279
280        // Safety: `av_write_frame` is safe to call, once the header has been written.
281        FfmpegErrorCode(unsafe { av_write_frame(self.as_mut_ptr(), packet.as_ptr() as *mut _) }).result()?;
282        Ok(())
283    }
284
285    /// Returns the flags for the output.
286    pub const fn flags(&self) -> AVFmtFlags {
287        AVFmtFlags(self.inner.context.as_deref_except().flags)
288    }
289
290    /// Returns the flags for the output.
291    pub const fn output_flags(&self) -> Option<AVFormatFlags> {
292        // Safety: The format is valid.
293        let Some(oformat) = (unsafe { self.inner.context.as_deref_except().oformat.as_ref() }) else {
294            return None;
295        };
296
297        Some(AVFormatFlags(oformat.flags))
298    }
299}
300
301impl Output<()> {
302    /// Opens the output with the given path.
303    pub fn open(path: &str) -> Result<Self, FfmpegError> {
304        Ok(Self {
305            inner: Inner::open_output(path)?,
306            state: OutputState::Uninitialized,
307        })
308    }
309}
310
311#[cfg(test)]
312#[cfg_attr(all(test, coverage_nightly), coverage(off))]
313mod tests {
314    use std::ffi::CString;
315    use std::io::{Cursor, Write};
316    use std::ptr;
317
318    use bytes::{Buf, Bytes};
319    use sha2::Digest;
320    use tempfile::Builder;
321
322    use crate::dict::Dictionary;
323    use crate::error::FfmpegError;
324    use crate::io::output::{AVCodec, AVRational, OutputState};
325    use crate::io::{Input, Output, OutputOptions};
326    use crate::{AVFmtFlags, AVMediaType, file_path};
327
328    #[test]
329    fn test_output_options_get_format_ffi_null() {
330        let format_name = CString::new("mp4").unwrap();
331        let format_mime_type = CString::new("").unwrap();
332        // Safety: `av_guess_format` is safe to call and all arguments are valid.
333        let format_ptr =
334            unsafe { crate::ffi::av_guess_format(format_name.as_ptr(), ptr::null(), format_mime_type.as_ptr()) };
335
336        assert!(
337            !format_ptr.is_null(),
338            "Failed to retrieve AVOutputFormat for the given format name"
339        );
340
341        let output_options = OutputOptions::builder().format_name("mp4").unwrap().build();
342        assert_eq!(output_options.format_ffi, format_ptr);
343    }
344
345    #[test]
346    fn test_output_options_get_format_ffi_output_format_error() {
347        match OutputOptions::builder().format_name("unknown_format") {
348            Ok(_) => panic!("Expected error, got Ok"),
349            Err(e) => {
350                assert_eq!(e, FfmpegError::Arguments("could not determine output format"));
351            }
352        }
353    }
354
355    #[test]
356    fn test_output_into_inner() {
357        let data = Cursor::new(Vec::with_capacity(1024));
358        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
359        let output = Output::new(data, options).expect("Failed to create Output");
360        let inner_data = output.into_inner();
361
362        assert!(inner_data.get_ref().is_empty());
363        let buffer = inner_data.into_inner();
364        assert_eq!(buffer.capacity(), 1024);
365    }
366
367    #[test]
368    fn test_output_new() {
369        let data = Cursor::new(Vec::new());
370        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
371        let output = Output::new(data, options);
372
373        assert!(output.is_ok());
374    }
375
376    #[test]
377    fn test_output_seekable() {
378        let data = Cursor::new(Vec::new());
379        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
380        let output = Output::seekable(data, options);
381
382        assert!(output.is_ok());
383    }
384
385    #[test]
386    fn test_output_set_metadata() {
387        let data = Cursor::new(Vec::new());
388        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
389        let mut output = Output::new(data, options).unwrap();
390        let metadata = Dictionary::new();
391        output.set_metadata(metadata);
392
393        assert!(!output.as_ptr().is_null());
394    }
395
396    #[test]
397    fn test_output_as_mut_ptr() {
398        let data = Cursor::new(Vec::new());
399        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
400        let mut output = Output::new(data, options).expect("Failed to create Output");
401        let context_ptr = output.as_mut_ptr();
402
403        assert!(!context_ptr.is_null(), "Expected non-null pointer from as_mut_ptr");
404    }
405
406    #[test]
407    fn test_add_stream_with_valid_codec() {
408        let data = Cursor::new(Vec::new());
409        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
410        let mut output = Output::new(data, options).expect("Failed to create Output");
411        let dummy_codec: *const AVCodec = 0x1234 as *const AVCodec;
412        let stream = output.add_stream(Some(dummy_codec));
413
414        assert!(stream.is_some(), "Expected a valid Stream to be added");
415    }
416
417    #[test]
418    fn test_copy_stream() {
419        let data = Cursor::new(Vec::new());
420        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
421        let mut output = Output::new(data, options).expect("Failed to create Output");
422
423        // create new output to prevent double mut borrow
424        let data = Cursor::new(Vec::new());
425        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
426        let mut output_two = Output::new(data, options).expect("Failed to create Output");
427
428        let dummy_codec: *const AVCodec = 0x1234 as *const AVCodec;
429        let mut source_stream = output.add_stream(Some(dummy_codec)).expect("Failed to add source stream");
430
431        source_stream.set_time_base(AVRational { num: 1, den: 25 });
432        source_stream.set_start_time(Some(1000));
433        source_stream.set_duration(Some(500));
434        let copied_stream = output_two
435            .copy_stream(&source_stream)
436            .expect("Failed to copy the stream")
437            .expect("Failed to copy the stream");
438
439        assert_eq!(copied_stream.index(), source_stream.index(), "Stream indices should match");
440        assert_eq!(copied_stream.id(), source_stream.id(), "Stream IDs should match");
441        assert_eq!(
442            copied_stream.time_base(),
443            source_stream.time_base(),
444            "Time bases should match"
445        );
446        assert_eq!(
447            copied_stream.start_time(),
448            source_stream.start_time(),
449            "Start times should match"
450        );
451        assert_eq!(copied_stream.duration(), source_stream.duration(), "Durations should match");
452        assert_eq!(copied_stream.duration(), source_stream.duration(), "Durations should match");
453        assert!(!copied_stream.as_ptr().is_null(), "Copied stream pointer should not be null");
454    }
455
456    #[test]
457    fn test_output_flags() {
458        let data = Cursor::new(Vec::new());
459        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
460        let output = Output::new(data, options).expect("Failed to create Output");
461        let flags = output.flags();
462
463        assert_eq!(flags, AVFmtFlags::AutoBsf, "Expected default flag to be AVFMT_FLAG_AUTO_BSF");
464    }
465
466    #[test]
467    fn test_output_open() {
468        let temp_file = Builder::new()
469            .suffix(".mp4")
470            .tempfile()
471            .expect("Failed to create a temporary file");
472        let temp_path = temp_file.path();
473        let output = Output::open(temp_path.to_str().unwrap());
474
475        assert!(output.is_ok(), "Expected Output::open to succeed");
476    }
477
478    macro_rules! get_boxes {
479        ($output:expr) => {{
480            let binary = $output.inner.data.as_mut().unwrap().get_mut().as_slice();
481            let mut cursor = Cursor::new(Bytes::copy_from_slice(binary));
482            let mut boxes = Vec::new();
483            while cursor.has_remaining() {
484                let mut box_ = scuffle_mp4::DynBox::demux(&mut cursor).expect("Failed to demux mp4");
485                if let scuffle_mp4::DynBox::Mdat(mdat) = &mut box_ {
486                    mdat.data.iter_mut().for_each(|buf| {
487                        let mut hash = sha2::Sha256::new();
488                        hash.write_all(buf).unwrap();
489                        *buf = hash.finalize().to_vec().into();
490                    });
491                }
492                boxes.push(box_);
493            }
494
495            boxes
496        }};
497    }
498
499    #[test]
500    fn test_output_write_mp4() {
501        let data = Cursor::new(Vec::new());
502        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
503
504        let mut output = Output::seekable(data, options).expect("Failed to create Output");
505
506        let mut input = Input::seekable(std::fs::File::open(file_path("avc_aac.mp4")).expect("Failed to open file"))
507            .expect("Failed to create Input");
508        let streams = input.streams();
509        let best_video_stream = streams.best(AVMediaType::Video).expect("no video stream found");
510
511        output.copy_stream(&best_video_stream).expect("Failed to copy stream");
512
513        output.write_header().expect("Failed to write header");
514        assert_eq!(output.state, OutputState::HeaderWritten, "Expected header to be written");
515        assert!(output.write_header().is_err(), "Expected error when writing header twice");
516
517        insta::assert_debug_snapshot!("test_output_write_mp4_header", get_boxes!(output));
518
519        let best_video_stream_index = best_video_stream.index();
520
521        while let Some(packet) = input.receive_packet().expect("Failed to receive packet") {
522            if packet.stream_index() != best_video_stream_index {
523                continue;
524            }
525
526            output.write_interleaved_packet(packet).expect("Failed to write packet");
527        }
528
529        insta::assert_debug_snapshot!("test_output_write_mp4_packets", get_boxes!(output));
530
531        output.write_trailer().expect("Failed to write trailer");
532        assert!(output.write_trailer().is_err(), "Expected error when writing trailer twice");
533        assert_eq!(output.state, OutputState::TrailerWritten, "Expected trailer to be written");
534
535        insta::assert_debug_snapshot!("test_output_write_mp4_trailer", get_boxes!(output));
536    }
537
538    #[test]
539    fn test_output_write_mp4_fragmented() {
540        let data = Cursor::new(Vec::new());
541        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
542
543        let mut output = Output::seekable(data, options).expect("Failed to create Output");
544
545        let mut input = Input::seekable(std::fs::File::open(file_path("avc_aac.mp4")).expect("Failed to open file"))
546            .expect("Failed to create Input");
547        let streams = input.streams();
548        let best_video_stream = streams.best(AVMediaType::Video).expect("no video stream found");
549
550        output.copy_stream(&best_video_stream).expect("Failed to copy stream");
551
552        output
553            .write_header_with_options(
554                &mut Dictionary::try_from_iter([("movflags", "frag_keyframe+empty_moov")])
555                    .expect("Failed to create dictionary from hashmap"),
556            )
557            .expect("Failed to write header");
558        assert_eq!(output.state, OutputState::HeaderWritten, "Expected header to be written");
559        assert!(
560            output
561                .write_header_with_options(
562                    &mut Dictionary::try_from_iter([("movflags", "frag_keyframe+empty_moov")],)
563                        .expect("Failed to create dictionary from hashmap")
564                )
565                .is_err(),
566            "Expected error when writing header twice"
567        );
568
569        insta::assert_debug_snapshot!("test_output_write_mp4_fragmented_header", get_boxes!(output));
570
571        let best_video_stream_index = best_video_stream.index();
572
573        while let Some(packet) = input.receive_packet().expect("Failed to receive packet") {
574            if packet.stream_index() != best_video_stream_index {
575                continue;
576            }
577
578            output.write_packet(&packet).expect("Failed to write packet");
579        }
580
581        insta::assert_debug_snapshot!("test_output_write_mp4_fragmented_packets", get_boxes!(output));
582
583        output.write_trailer().expect("Failed to write trailer");
584        assert_eq!(output.state, OutputState::TrailerWritten, "Expected trailer to be written");
585        assert!(output.write_trailer().is_err(), "Expected error when writing trailer twice");
586
587        insta::assert_debug_snapshot!("test_output_write_mp4_fragmented_trailer", get_boxes!(output));
588    }
589}
scuffle_ffmpeg/io/output.rs

scuffle_ffmpeg/io/
output.rs
