scuffle_ffmpeg/io/
output.rs

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

scuffle_ffmpeg/io/
output.rs
