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::path::PathBuf;
317    use std::ptr;
318
319    use bytes::{Buf, Bytes};
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};
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 cursor = Cursor::new(Bytes::copy_from_slice(binary));
483            let mut boxes = Vec::new();
484            while cursor.has_remaining() {
485                let mut box_ = scuffle_mp4::DynBox::demux(&mut cursor).expect("Failed to demux mp4");
486                if let scuffle_mp4::DynBox::Mdat(mdat) = &mut box_ {
487                    mdat.data.iter_mut().for_each(|buf| {
488                        let mut hash = sha2::Sha256::new();
489                        hash.write_all(buf).unwrap();
490                        *buf = hash.finalize().to_vec().into();
491                    });
492                }
493                boxes.push(box_);
494            }
495
496            boxes
497        }};
498    }
499
500    #[test]
501    fn test_output_write_mp4() {
502        let data = Cursor::new(Vec::new());
503        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
504
505        let mut output = Output::seekable(data, options).expect("Failed to create Output");
506        let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../assets");
507
508        let mut input = Input::seekable(std::fs::File::open(dir.join("avc_aac.mp4")).expect("Failed to open file"))
509            .expect("Failed to create Input");
510        let streams = input.streams();
511        let best_video_stream = streams.best(AVMediaType::Video).expect("no video stream found");
512
513        output.copy_stream(&best_video_stream).expect("Failed to copy stream");
514
515        output.write_header().expect("Failed to write header");
516        assert_eq!(output.state, OutputState::HeaderWritten, "Expected header to be written");
517        assert!(output.write_header().is_err(), "Expected error when writing header twice");
518
519        insta::assert_debug_snapshot!("test_output_write_mp4_header", get_boxes!(output));
520
521        let best_video_stream_index = best_video_stream.index();
522
523        while let Some(packet) = input.receive_packet().expect("Failed to receive packet") {
524            if packet.stream_index() != best_video_stream_index {
525                continue;
526            }
527
528            output.write_interleaved_packet(packet).expect("Failed to write packet");
529        }
530
531        insta::assert_debug_snapshot!("test_output_write_mp4_packets", get_boxes!(output));
532
533        output.write_trailer().expect("Failed to write trailer");
534        assert!(output.write_trailer().is_err(), "Expected error when writing trailer twice");
535        assert_eq!(output.state, OutputState::TrailerWritten, "Expected trailer to be written");
536
537        insta::assert_debug_snapshot!("test_output_write_mp4_trailer", get_boxes!(output));
538    }
539
540    #[test]
541    fn test_output_write_mp4_fragmented() {
542        let data = Cursor::new(Vec::new());
543        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
544
545        let mut output = Output::seekable(data, options).expect("Failed to create Output");
546        let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../assets");
547
548        let mut input = Input::seekable(std::fs::File::open(dir.join("avc_aac.mp4")).expect("Failed to open file"))
549            .expect("Failed to create Input");
550        let streams = input.streams();
551        let best_video_stream = streams.best(AVMediaType::Video).expect("no video stream found");
552
553        output.copy_stream(&best_video_stream).expect("Failed to copy stream");
554
555        output
556            .write_header_with_options(
557                &mut Dictionary::try_from_iter([("movflags", "frag_keyframe+empty_moov")])
558                    .expect("Failed to create dictionary from hashmap"),
559            )
560            .expect("Failed to write header");
561        assert_eq!(output.state, OutputState::HeaderWritten, "Expected header to be written");
562        assert!(
563            output
564                .write_header_with_options(
565                    &mut Dictionary::try_from_iter([("movflags", "frag_keyframe+empty_moov")],)
566                        .expect("Failed to create dictionary from hashmap")
567                )
568                .is_err(),
569            "Expected error when writing header twice"
570        );
571
572        insta::assert_debug_snapshot!("test_output_write_mp4_fragmented_header", get_boxes!(output));
573
574        let best_video_stream_index = best_video_stream.index();
575
576        while let Some(packet) = input.receive_packet().expect("Failed to receive packet") {
577            if packet.stream_index() != best_video_stream_index {
578                continue;
579            }
580
581            output.write_packet(&packet).expect("Failed to write packet");
582        }
583
584        insta::assert_debug_snapshot!("test_output_write_mp4_fragmented_packets", get_boxes!(output));
585
586        output.write_trailer().expect("Failed to write trailer");
587        assert_eq!(output.state, OutputState::TrailerWritten, "Expected trailer to be written");
588        assert!(output.write_trailer().is_err(), "Expected error when writing trailer twice");
589
590        insta::assert_debug_snapshot!("test_output_write_mp4_fragmented_trailer", get_boxes!(output));
591    }
592}
scuffle_ffmpeg/io/output.rs

scuffle_ffmpeg/io/
output.rs
