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 bytes::{Buf, Bytes};
321    use sha2::Digest;
322    use tempfile::Builder;
323
324    use crate::dict::Dictionary;
325    use crate::error::FfmpegError;
326    use crate::io::output::{AVCodec, AVRational, OutputState};
327    use crate::io::{Input, Output, OutputOptions};
328    use crate::{AVFmtFlags, AVMediaType, file_path};
329
330    #[test]
331    fn test_output_options_get_format_ffi_null() {
332        let format_name = CString::new("mp4").unwrap();
333        let format_mime_type = CString::new("").unwrap();
334        // Safety: `av_guess_format` is safe to call and all arguments are valid.
335        let format_ptr =
336            unsafe { crate::ffi::av_guess_format(format_name.as_ptr(), ptr::null(), format_mime_type.as_ptr()) };
337
338        assert!(
339            !format_ptr.is_null(),
340            "Failed to retrieve AVOutputFormat for the given format name"
341        );
342
343        let output_options = OutputOptions::builder().format_name("mp4").unwrap().build();
344        assert_eq!(output_options.format_ffi, format_ptr);
345    }
346
347    #[test]
348    fn test_output_options_get_format_ffi_output_format_error() {
349        match OutputOptions::builder().format_name("unknown_format") {
350            Ok(_) => panic!("Expected error, got Ok"),
351            Err(e) => {
352                assert_eq!(e, FfmpegError::Arguments("could not determine output format"));
353            }
354        }
355    }
356
357    #[test]
358    fn test_output_into_inner() {
359        let data = Cursor::new(Vec::with_capacity(1024));
360        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
361        let output = Output::new(data, options).expect("Failed to create Output");
362        let inner_data = output.into_inner();
363
364        assert!(inner_data.get_ref().is_empty());
365        let buffer = inner_data.into_inner();
366        assert_eq!(buffer.capacity(), 1024);
367    }
368
369    #[test]
370    fn test_output_new() {
371        let data = Cursor::new(Vec::new());
372        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
373        let output = Output::new(data, options);
374
375        assert!(output.is_ok());
376    }
377
378    #[test]
379    fn test_output_seekable() {
380        let data = Cursor::new(Vec::new());
381        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
382        let output = Output::seekable(data, options);
383
384        assert!(output.is_ok());
385    }
386
387    #[test]
388    fn test_output_set_metadata() {
389        let data = Cursor::new(Vec::new());
390        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
391        let mut output = Output::new(data, options).unwrap();
392        let metadata = Dictionary::new();
393        output.set_metadata(metadata);
394
395        assert!(!output.as_ptr().is_null());
396    }
397
398    #[test]
399    fn test_output_as_mut_ptr() {
400        let data = Cursor::new(Vec::new());
401        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
402        let mut output = Output::new(data, options).expect("Failed to create Output");
403        let context_ptr = output.as_mut_ptr();
404
405        assert!(!context_ptr.is_null(), "Expected non-null pointer from as_mut_ptr");
406    }
407
408    #[test]
409    fn test_add_stream_with_valid_codec() {
410        let data = Cursor::new(Vec::new());
411        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
412        let mut output = Output::new(data, options).expect("Failed to create Output");
413        let dummy_codec: *const AVCodec = 0x1234 as *const AVCodec;
414        let stream = output.add_stream(Some(dummy_codec));
415
416        assert!(stream.is_some(), "Expected a valid Stream to be added");
417    }
418
419    #[test]
420    fn test_copy_stream() {
421        let data = Cursor::new(Vec::new());
422        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
423        let mut output = Output::new(data, options).expect("Failed to create Output");
424
425        // create new output to prevent double mut borrow
426        let data = Cursor::new(Vec::new());
427        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
428        let mut output_two = Output::new(data, options).expect("Failed to create Output");
429
430        let dummy_codec: *const AVCodec = 0x1234 as *const AVCodec;
431        let mut source_stream = output.add_stream(Some(dummy_codec)).expect("Failed to add source stream");
432
433        source_stream.set_time_base(AVRational { num: 1, den: 25 });
434        source_stream.set_start_time(Some(1000));
435        source_stream.set_duration(Some(500));
436        let copied_stream = output_two
437            .copy_stream(&source_stream)
438            .expect("Failed to copy the stream")
439            .expect("Failed to copy the stream");
440
441        assert_eq!(copied_stream.index(), source_stream.index(), "Stream indices should match");
442        assert_eq!(copied_stream.id(), source_stream.id(), "Stream IDs should match");
443        assert_eq!(
444            copied_stream.time_base(),
445            source_stream.time_base(),
446            "Time bases should match"
447        );
448        assert_eq!(
449            copied_stream.start_time(),
450            source_stream.start_time(),
451            "Start times should match"
452        );
453        assert_eq!(copied_stream.duration(), source_stream.duration(), "Durations should match");
454        assert_eq!(copied_stream.duration(), source_stream.duration(), "Durations should match");
455        assert!(!copied_stream.as_ptr().is_null(), "Copied stream pointer should not be null");
456    }
457
458    #[test]
459    fn test_output_flags() {
460        let data = Cursor::new(Vec::new());
461        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
462        let output = Output::new(data, options).expect("Failed to create Output");
463        let flags = output.flags();
464
465        assert_eq!(flags, AVFmtFlags::AutoBsf, "Expected default flag to be AVFMT_FLAG_AUTO_BSF");
466    }
467
468    #[test]
469    fn test_output_open() {
470        let temp_file = Builder::new()
471            .suffix(".mp4")
472            .tempfile()
473            .expect("Failed to create a temporary file");
474        let temp_path = temp_file.path();
475        let output = Output::open(temp_path.to_str().unwrap());
476
477        assert!(output.is_ok(), "Expected Output::open to succeed");
478    }
479
480    macro_rules! get_boxes {
481        ($output:expr) => {{
482            let binary = $output.inner.data.as_mut().unwrap().get_mut().as_slice();
483            let mut cursor = Cursor::new(Bytes::copy_from_slice(binary));
484            let mut boxes = Vec::new();
485            while cursor.has_remaining() {
486                let mut box_ = scuffle_mp4::DynBox::demux(&mut cursor).expect("Failed to demux mp4");
487                if let scuffle_mp4::DynBox::Mdat(mdat) = &mut box_ {
488                    mdat.data.iter_mut().for_each(|buf| {
489                        let mut hash = sha2::Sha256::new();
490                        hash.write_all(buf).unwrap();
491                        *buf = hash.finalize().to_vec().into();
492                    });
493                }
494                boxes.push(box_);
495            }
496
497            boxes
498        }};
499    }
500
501    #[test]
502    fn test_output_write_mp4() {
503        let data = Cursor::new(Vec::new());
504        let options = OutputOptions::builder().format_name("mp4").unwrap().build();
505
506        let mut output = Output::seekable(data, options).expect("Failed to create Output");
507
508        let mut input = Input::seekable(std::fs::File::open(file_path("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
547        let mut input = Input::seekable(std::fs::File::open(file_path("avc_aac.mp4")).expect("Failed to open file"))
548            .expect("Failed to create Input");
549        let streams = input.streams();
550        let best_video_stream = streams.best(AVMediaType::Video).expect("no video stream found");
551
552        output.copy_stream(&best_video_stream).expect("Failed to copy stream");
553
554        output
555            .write_header_with_options(
556                &mut Dictionary::try_from_iter([("movflags", "frag_keyframe+empty_moov")])
557                    .expect("Failed to create dictionary from hashmap"),
558            )
559            .expect("Failed to write header");
560        assert_eq!(output.state, OutputState::HeaderWritten, "Expected header to be written");
561        assert!(
562            output
563                .write_header_with_options(
564                    &mut Dictionary::try_from_iter([("movflags", "frag_keyframe+empty_moov")],)
565                        .expect("Failed to create dictionary from hashmap")
566                )
567                .is_err(),
568            "Expected error when writing header twice"
569        );
570
571        insta::assert_debug_snapshot!("test_output_write_mp4_fragmented_header", get_boxes!(output));
572
573        let best_video_stream_index = best_video_stream.index();
574
575        while let Some(packet) = input.receive_packet().expect("Failed to receive packet") {
576            if packet.stream_index() != best_video_stream_index {
577                continue;
578            }
579
580            output.write_packet(&packet).expect("Failed to write packet");
581        }
582
583        insta::assert_debug_snapshot!("test_output_write_mp4_fragmented_packets", get_boxes!(output));
584
585        output.write_trailer().expect("Failed to write trailer");
586        assert_eq!(output.state, OutputState::TrailerWritten, "Expected trailer to be written");
587        assert!(output.write_trailer().is_err(), "Expected error when writing trailer twice");
588
589        insta::assert_debug_snapshot!("test_output_write_mp4_fragmented_trailer", get_boxes!(output));
590    }
591}
scuffle_ffmpeg/io/output.rs

scuffle_ffmpeg/io/
output.rs
