scuffle_ffmpeg/io/
internal.rs

1use libc::c_void;
2
3use crate::error::{FfmpegError, FfmpegErrorCode};
4use crate::ffi::*;
5use crate::smart_object::SmartPtr;
6use crate::{AVIOFlag, AVSeekWhence};
7
8const AVERROR_IO: i32 = AVERROR(EIO);
9
10/// Safety: The function must be used with the same type as the one used to
11/// generically create the function pointer
12pub(crate) unsafe extern "C" fn read_packet<T: std::io::Read>(
13    opaque: *mut libc::c_void,
14    buf: *mut u8,
15    buf_size: i32,
16) -> i32 {
17    // Safety: The pointer is valid given the way this function is constructed, the opaque pointer is a pointer to a T.
18    let this = unsafe { &mut *(opaque as *mut T) };
19    // Safety: the buffer has at least `buf_size` bytes.
20    let buffer = unsafe { std::slice::from_raw_parts_mut(buf, buf_size as usize) };
21
22    let ret = this.read(buffer).map(|n| n as i32).unwrap_or(AVERROR_IO);
23
24    if ret == 0 {
25        return AVERROR_EOF;
26    }
27
28    ret
29}
30
31/// Safety: The function must be used with the same type as the one used to
32/// generically create the function pointer
33pub(crate) unsafe extern "C" fn write_packet<T: std::io::Write>(
34    opaque: *mut libc::c_void,
35    buf: *const u8,
36    buf_size: i32,
37) -> i32 {
38    // Safety: The pointer is valid given the way this function is constructed, the opaque pointer is a pointer to a T.
39    let this = unsafe { &mut *(opaque as *mut T) };
40    // Safety: the buffer has at least `buf_size` bytes.
41    let buffer = unsafe { std::slice::from_raw_parts(buf, buf_size as usize) };
42
43    this.write(buffer).map(|n| n as i32).unwrap_or(AVERROR_IO)
44}
45
46/// Safety: The function must be used with the same type as the one used to
47/// generically create the function pointer
48pub(crate) unsafe extern "C" fn seek<T: std::io::Seek>(opaque: *mut libc::c_void, offset: i64, whence: i32) -> i64 {
49    // Safety: The pointer is valid given the way this function is constructed, the opaque pointer is a pointer to a T.
50    let this = unsafe { &mut *(opaque as *mut T) };
51
52    let mut whence = AVSeekWhence(whence);
53
54    let seek_size = whence & AVSeekWhence::Size != 0;
55    if seek_size {
56        whence &= !AVSeekWhence::Size;
57    }
58
59    let seek_force = whence & AVSeekWhence::Force != 0;
60    if seek_force {
61        whence &= !AVSeekWhence::Force;
62    }
63
64    if seek_size {
65        let Ok(pos) = this.stream_position() else {
66            return AVERROR_IO as i64;
67        };
68
69        let Ok(end) = this.seek(std::io::SeekFrom::End(0)) else {
70            return AVERROR_IO as i64;
71        };
72
73        if end != pos {
74            let Ok(_) = this.seek(std::io::SeekFrom::Start(pos)) else {
75                return AVERROR_IO as i64;
76            };
77        }
78
79        return end as i64;
80    }
81
82    let whence = match whence {
83        AVSeekWhence::Start => std::io::SeekFrom::Start(offset as u64),
84        AVSeekWhence::Current => std::io::SeekFrom::Current(offset),
85        AVSeekWhence::End => std::io::SeekFrom::End(offset),
86        _ => return -1,
87    };
88
89    match this.seek(whence) {
90        Ok(pos) => pos as i64,
91        Err(_) => AVERROR_IO as i64,
92    }
93}
94
95pub(crate) struct Inner<T: Send + Sync> {
96    pub(crate) data: Option<Box<T>>,
97    pub(crate) context: SmartPtr<AVFormatContext>,
98    _io: SmartPtr<AVIOContext>,
99}
100
101pub(crate) struct InnerOptions {
102    pub(crate) buffer_size: usize,
103    pub(crate) read_fn: Option<unsafe extern "C" fn(*mut c_void, *mut u8, i32) -> i32>,
104    pub(crate) write_fn: Option<unsafe extern "C" fn(*mut c_void, *const u8, i32) -> i32>,
105    pub(crate) seek_fn: Option<unsafe extern "C" fn(*mut c_void, i64, i32) -> i64>,
106    pub(crate) output_format: *const AVOutputFormat,
107}
108
109impl Default for InnerOptions {
110    fn default() -> Self {
111        Self {
112            buffer_size: 4096,
113            read_fn: None,
114            write_fn: None,
115            seek_fn: None,
116            output_format: std::ptr::null(),
117        }
118    }
119}
120
121impl<T: Send + Sync> Inner<T> {
122    /// Creates a new `Inner` instance.
123    pub(crate) fn new(data: T, options: InnerOptions) -> Result<Self, FfmpegError> {
124        // Safety: av_malloc is safe to call
125        let buffer = unsafe { av_malloc(options.buffer_size) };
126
127        fn buffer_destructor(ptr: &mut *mut c_void) {
128            // We own this resource so we need to free it
129            // Safety: this buffer was allocated via `av_malloc` so we need to free it.
130            unsafe { av_free(*ptr) };
131            // We clear the old pointer so it doesn't get freed again.
132            *ptr = std::ptr::null_mut();
133        }
134
135        // This is only a temporary smart_ptr because we will change the ownership to be owned by the io context. & freed, by the io context.
136        // Safety: av_malloc gives a valid pointer & the destructor has been setup to free the buffer.
137        let buffer = unsafe { SmartPtr::wrap_non_null(buffer, buffer_destructor) }.ok_or(FfmpegError::Alloc)?;
138
139        let mut data = Box::new(data);
140
141        // Safety: avio_alloc_context is safe to call, and all the function pointers are valid
142        let destructor = |ptr: &mut *mut AVIOContext| {
143            // Safety: the pointer is always valid.
144            let mut_ref = unsafe { ptr.as_mut() };
145            if let Some(ptr) = mut_ref {
146                buffer_destructor(&mut (ptr.buffer as *mut c_void));
147            }
148
149            // Safety: avio_context_free is safe to call & the pointer was allocated by `av_create_io_context`
150            unsafe { avio_context_free(ptr) };
151            *ptr = std::ptr::null_mut();
152        };
153
154        // Safety: avio_alloc_context is safe to call & the destructor has been setup to free the buffer.
155        let io = unsafe {
156            avio_alloc_context(
157                buffer.as_ptr() as *mut u8,
158                options.buffer_size as i32,
159                if options.write_fn.is_some() { 1 } else { 0 },
160                data.as_mut() as *mut _ as *mut c_void,
161                options.read_fn,
162                options.write_fn,
163                options.seek_fn,
164            )
165        };
166
167        // Safety: `avio_alloc_context` returns a valid pointer & the destructor has been setup to free both the buffer & the io context.
168        let mut io = unsafe { SmartPtr::wrap_non_null(io, destructor) }.ok_or(FfmpegError::Alloc)?;
169
170        // The buffer is now owned by the IO context. We need to go into_inner here to prevent the destructor from being called before we use it.
171        buffer.into_inner();
172
173        let mut context = if options.write_fn.is_some() {
174            let mut context = SmartPtr::null(|mut_ref| {
175                let ptr = *mut_ref;
176                // Safety: The pointer here is valid.
177                unsafe { avformat_free_context(ptr) };
178                *mut_ref = std::ptr::null_mut();
179            });
180
181            // Safety: avformat_alloc_output_context2 is safe to call
182            FfmpegErrorCode(unsafe {
183                avformat_alloc_output_context2(
184                    context.as_mut(),
185                    options.output_format,
186                    std::ptr::null(),
187                    std::ptr::null_mut(),
188                )
189            })
190            .result()?;
191
192            if context.as_ptr().is_null() {
193                return Err(FfmpegError::Alloc);
194            }
195
196            context
197        } else {
198            // Safety: avformat_alloc_context is safe to call
199            let context = unsafe { avformat_alloc_context() };
200
201            let destructor = |mut_ref: &mut *mut AVFormatContext| {
202                let ptr = *mut_ref;
203                // Safety: The pointer here is valid and was allocated by `avformat_alloc_context`.
204                unsafe { avformat_free_context(ptr) };
205                *mut_ref = std::ptr::null_mut();
206            };
207
208            // Safety: `avformat_alloc_context` returns a valid pointer & the destructor has been setup to free the context.
209            unsafe { SmartPtr::wrap_non_null(context, destructor) }.ok_or(FfmpegError::Alloc)?
210        };
211
212        // The io context will live as long as the format context
213        context.as_deref_mut().expect("Context is null").pb = io.as_mut_ptr();
214
215        Ok(Self {
216            data: Some(data),
217            context,
218            _io: io,
219        })
220    }
221}
222
223impl Inner<()> {
224    /// Empty context cannot be used until its initialized and setup correctly
225    /// Safety: this function is marked as unsafe because it must be initialized and setup correctltly before returning it to the user.
226    pub(crate) unsafe fn empty() -> Self {
227        Self {
228            data: Some(Box::new(())),
229            context: SmartPtr::null(|mut_ref| {
230                // We own this resource so we need to free it
231                let ptr = *mut_ref;
232                // Safety: The pointer here is valid.
233                unsafe { avformat_free_context(ptr) };
234                *mut_ref = std::ptr::null_mut();
235            }),
236            _io: SmartPtr::null(|_| {}),
237        }
238    }
239
240    /// Opens an output stream to a file path.
241    pub(crate) fn open_output(path: &str) -> Result<Self, FfmpegError> {
242        let path = std::ffi::CString::new(path).expect("Failed to convert path to CString");
243
244        // Safety: We immediately initialize the inner and setup the context.
245        let mut this = unsafe { Self::empty() };
246
247        // Safety: avformat_alloc_output_context2 is safe to call and all arguments are valid
248        FfmpegErrorCode(unsafe {
249            avformat_alloc_output_context2(this.context.as_mut(), std::ptr::null(), std::ptr::null(), path.as_ptr())
250        })
251        .result()?;
252
253        // We are not moving the pointer so this is safe
254        if this.context.as_ptr().is_null() {
255            return Err(FfmpegError::Alloc);
256        }
257
258        // Safety: avio_open is safe to call and all arguments are valid
259        FfmpegErrorCode(unsafe {
260            avio_open(
261                &mut this.context.as_deref_mut_except().pb,
262                path.as_ptr(),
263                AVIOFlag::Write.into(),
264            )
265        })
266        .result()?;
267
268        this.context.set_destructor(|mut_ref| {
269            // We own this resource so we need to free it
270            let ptr = *mut_ref;
271            // Safety: The pointer here is valid.
272            let mut_ptr_ref = unsafe { ptr.as_mut() };
273
274            if let Some(mut_ptr_ref) = mut_ptr_ref {
275                // Safety: The pointer here is valid. And we need to clean this up before we do `avformat_free_context`.
276                unsafe { avio_closep(&mut mut_ptr_ref.pb) };
277            }
278
279            // Safety: The pointer here is valid.
280            unsafe { avformat_free_context(ptr) };
281            *mut_ref = std::ptr::null_mut();
282        });
283
284        Ok(this)
285    }
286}
287
288#[cfg(test)]
289#[cfg_attr(all(test, coverage_nightly), coverage(off))]
290mod tests {
291    use std::ffi::CString;
292    use std::io::Cursor;
293    use std::sync::Once;
294    use std::sync::atomic::{AtomicUsize, Ordering};
295
296    use libc::c_void;
297    use tempfile::Builder;
298
299    use crate::AVSeekWhence;
300    use crate::error::FfmpegError;
301    use crate::ffi::av_guess_format;
302    use crate::io::internal::{AVERROR_EOF, Inner, InnerOptions, read_packet, seek, write_packet};
303
304    #[test]
305    fn test_read_packet_eof() {
306        let mut data: Cursor<Vec<u8>> = Cursor::new(vec![]);
307        let mut buf = [0u8; 10];
308
309        // Safety: The pointer is valid.
310        unsafe {
311            let result =
312                read_packet::<Cursor<Vec<u8>>>((&raw mut data) as *mut libc::c_void, buf.as_mut_ptr(), buf.len() as i32);
313
314            assert_eq!(result, AVERROR_EOF);
315        }
316    }
317
318    #[test]
319    fn test_write_packet_success() {
320        let mut data = Cursor::new(vec![0u8; 10]);
321        let buf = [1u8, 2, 3, 4, 5];
322
323        // Safety: The pointer is valid.
324        unsafe {
325            let result = write_packet::<Cursor<Vec<u8>>>((&raw mut data) as *mut c_void, buf.as_ptr(), buf.len() as i32);
326            assert_eq!(result, buf.len() as i32);
327
328            let written_data = data.get_ref();
329            assert_eq!(&written_data[..buf.len()], &buf);
330        }
331    }
332
333    #[test]
334    fn test_seek_force() {
335        let mut cursor = Cursor::new(vec![0u8; 100]);
336        let opaque = &raw mut cursor as *mut c_void;
337        assert_eq!(cursor.position(), 0);
338        let offset = 10;
339        let mut whence = AVSeekWhence::Current | AVSeekWhence::Force;
340        // Safety: The pointer is valid.
341        let result = unsafe { seek::<Cursor<Vec<u8>>>(opaque, offset, whence.into()) };
342
343        assert_eq!(result, { offset });
344        whence &= !AVSeekWhence::Force;
345        assert_eq!(whence, AVSeekWhence::Current);
346        assert_eq!(cursor.position(), offset as u64);
347    }
348
349    #[test]
350    fn test_seek_seek_end() {
351        let mut cursor = Cursor::new(vec![0u8; 100]);
352        let opaque = &raw mut cursor as *mut libc::c_void;
353        let offset = -10;
354        // Safety: The pointer is valid.
355        let result = unsafe { seek::<Cursor<Vec<u8>>>(opaque, offset, AVSeekWhence::End.into()) };
356
357        assert_eq!(result, 90);
358        assert_eq!(cursor.position(), 90);
359    }
360
361    #[test]
362    fn test_seek_invalid_whence() {
363        let mut cursor = Cursor::new(vec![0u8; 100]);
364        let opaque = &raw mut cursor as *mut libc::c_void;
365        // Safety: The pointer is valid.
366        let result = unsafe { seek::<Cursor<Vec<u8>>>(opaque, 0, 999) };
367
368        assert_eq!(result, -1);
369        assert_eq!(cursor.position(), 0);
370    }
371
372    #[test]
373    fn test_avformat_alloc_output_context2_error() {
374        static BUF_SIZE_TRACKER: AtomicUsize = AtomicUsize::new(0);
375        static CALL_COUNT: AtomicUsize = AtomicUsize::new(0);
376        static INIT: Once = Once::new();
377
378        INIT.call_once(|| {
379            BUF_SIZE_TRACKER.store(0, Ordering::SeqCst);
380            CALL_COUNT.store(0, Ordering::SeqCst);
381        });
382
383        unsafe extern "C" fn dummy_write_fn(_opaque: *mut libc::c_void, _buf: *const u8, _buf_size: i32) -> i32 {
384            CALL_COUNT.fetch_add(1, Ordering::SeqCst);
385            BUF_SIZE_TRACKER.store(_buf_size as usize, Ordering::SeqCst);
386            0 // simulate success
387        }
388
389        let invalid_format = CString::new("invalid_format").expect("Failed to create CString");
390        let options = InnerOptions {
391            buffer_size: 4096,
392            write_fn: Some(dummy_write_fn),
393            // Safety: av_guess_format is safe to call
394            output_format: unsafe { av_guess_format(invalid_format.as_ptr(), std::ptr::null(), std::ptr::null()) },
395            ..Default::default()
396        };
397        let data = ();
398        let result = Inner::new(data, options);
399
400        assert!(result.is_err(), "Expected an error but got Ok");
401
402        let call_count = CALL_COUNT.load(Ordering::SeqCst);
403        assert_eq!(call_count, 0, "Expected dummy_write_fn to not be called.");
404
405        if let Err(error) = result {
406            match error {
407                FfmpegError::Code(_) => {
408                    eprintln!("Expected avformat_alloc_output_context2 error occurred.");
409                }
410                _ => panic!("Unexpected error variant: {error:?}"),
411            }
412        }
413    }
414
415    #[test]
416    fn test_open_output_valid_path() {
417        let temp_file = Builder::new()
418            .suffix(".mp4")
419            .tempfile()
420            .expect("Failed to create a temporary file");
421        let test_path = temp_file.path();
422        let result = Inner::open_output(test_path.to_str().unwrap());
423
424        assert!(result.is_ok(), "Expected success but got error");
425    }
426
427    #[test]
428    fn test_open_output_invalid_path() {
429        let test_path = "";
430        let result = Inner::open_output(test_path);
431
432        assert!(result.is_err(), "Expected Err, got Ok");
433    }
434
435    #[test]
436    fn test_open_output_avformat_alloc_error() {
437        let test_path = tempfile::tempdir().unwrap().path().join("restricted_output.mp4");
438        let test_path_str = test_path.to_str().unwrap();
439        let result = Inner::open_output(test_path_str);
440        if let Err(error) = &result {
441            eprintln!("Function returned an error: {error:?}");
442        }
443
444        assert!(
445            matches!(result, Err(FfmpegError::Code(_))),
446            "Expected FfmpegError::Code but received a different error."
447        );
448    }
449}