xref: /aosp_15_r20/external/crosvm/media/ffmpeg/src/avcodec.rs (revision bb4ee6a4ae7042d18b07a98463b9c8b875e44b39)
1 // Copyright 2022 The ChromiumOS Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 //! This module implements a lightweight and safe decoder interface over `libavcodec`. It is
6 //! designed to concentrate all calls to unsafe methods in one place, while providing the same
7 //! low-level access as the libavcodec functions do.
8 
9 use std::ffi::CStr;
10 use std::fmt::Debug;
11 use std::fmt::Display;
12 use std::marker::PhantomData;
13 use std::mem::ManuallyDrop;
14 use std::ops::Deref;
15 
16 use libc::c_char;
17 use libc::c_int;
18 use libc::c_void;
19 use thiserror::Error as ThisError;
20 
21 use super::*;
22 use crate::ffi::AVPictureType;
23 
24 /// An error returned by a low-level libavcodec function.
25 #[derive(Debug, ThisError)]
26 pub struct AvError(pub libc::c_int);
27 
28 impl AvError {
result(ret: c_int) -> Result<(), Self>29     pub fn result(ret: c_int) -> Result<(), Self> {
30         if ret >= 0 {
31             Ok(())
32         } else {
33             Err(AvError(ret))
34         }
35     }
36 }
37 
38 impl Display for AvError {
fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result39     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40         let mut buffer = [0u8; 255];
41         let ret =
42             // SAFETY:
43             // Safe because we are passing valid bounds for the buffer.
44             unsafe { ffi::av_strerror(self.0, buffer.as_mut_ptr() as *mut c_char, buffer.len()) };
45         match ret {
46             ret if ret >= 0 => {
47                 let end_of_string = buffer.iter().position(|i| *i == 0).unwrap_or(buffer.len());
48                 let error_string = std::string::String::from_utf8_lossy(&buffer[..end_of_string]);
49                 f.write_str(&error_string)
50             }
51             _ => f.write_fmt(format_args!("Unknown avcodec error {}", self.0)),
52         }
53     }
54 }
55 
56 /// Lightweight abstraction over libavcodec's `AVCodec` struct, allowing the query the capabilities
57 /// of supported codecs and opening a session to work with them.
58 ///
59 /// `AVCodec` instances in libavcodec are all static, hence we can safely use a static reference
60 /// lifetime here.
61 pub struct AvCodec(&'static ffi::AVCodec);
62 
63 #[derive(Debug, ThisError)]
64 pub enum AvCodecOpenError {
65     #[error("failed to allocate AVContext object")]
66     ContextAllocation,
67     #[error("failed to open AVContext object")]
68     ContextOpen,
69     #[error("ContextBuilder variant does not match codec type")]
70     UnexpectedCodecType,
71 }
72 
73 /// Dimensions of a frame, used in AvCodecContext and AvFrame.
74 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
75 pub struct Dimensions {
76     pub width: u32,
77     pub height: u32,
78 }
79 
80 impl AvCodec {
81     /// Returns whether the codec is a decoder.
is_decoder(&self) -> bool82     pub fn is_decoder(&self) -> bool {
83         // SAFETY:
84         // Safe because `av_codec_is_decoder` is called on a valid static `AVCodec` reference.
85         (unsafe { ffi::av_codec_is_decoder(self.0) } != 0)
86     }
87 
88     /// Returns whether the codec is an encoder.
is_encoder(&self) -> bool89     pub fn is_encoder(&self) -> bool {
90         // SAFETY:
91         // Safe because `av_codec_is_encoder` is called on a valid static `AVCodec` reference.
92         (unsafe { ffi::av_codec_is_encoder(self.0) } != 0)
93     }
94 
95     /// Returns the name of the codec.
name(&self) -> &'static str96     pub fn name(&self) -> &'static str {
97         const INVALID_CODEC_STR: &str = "invalid codec";
98 
99         // SAFETY:
100         // Safe because `CStr::from_ptr` is called on a valid zero-terminated C string.
101         unsafe { CStr::from_ptr(self.0.name).to_str() }.unwrap_or(INVALID_CODEC_STR)
102     }
103 
104     /// Returns the capabilities of the codec, as a mask of AV_CODEC_CAP_* bits.
capabilities(&self) -> u32105     pub fn capabilities(&self) -> u32 {
106         self.0.capabilities as u32
107     }
108 
109     /// Returns an iterator over the profiles supported by this codec.
profile_iter(&self) -> AvProfileIterator110     pub fn profile_iter(&self) -> AvProfileIterator {
111         AvProfileIterator(self.0.profiles)
112     }
113 
114     /// Returns an iterator over the pixel formats supported by this codec.
115     ///
116     /// For a decoder, the returned array will likely be empty. This means that ffmpeg's native
117     /// pixel format (YUV420) will be used.
pixel_format_iter(&self) -> AvPixelFormatIterator118     pub fn pixel_format_iter(&self) -> AvPixelFormatIterator {
119         AvPixelFormatIterator(self.0.pix_fmts)
120     }
121 
122     /// Get a builder for a encoder [`AvCodecContext`] using this codec.
build_encoder(&self) -> Result<EncoderContextBuilder, AvCodecOpenError>123     pub fn build_encoder(&self) -> Result<EncoderContextBuilder, AvCodecOpenError> {
124         if !self.is_encoder() {
125             return Err(AvCodecOpenError::UnexpectedCodecType);
126         }
127 
128         Ok(EncoderContextBuilder {
129             codec: self.0,
130             context: self.alloc_context()?,
131         })
132     }
133 
134     /// Get a builder for a decoder [`AvCodecContext`] using this codec.
build_decoder(&self) -> Result<DecoderContextBuilder, AvCodecOpenError>135     pub fn build_decoder(&self) -> Result<DecoderContextBuilder, AvCodecOpenError> {
136         if !self.is_decoder() {
137             return Err(AvCodecOpenError::UnexpectedCodecType);
138         }
139 
140         Ok(DecoderContextBuilder {
141             codec: self.0,
142             context: self.alloc_context()?,
143         })
144     }
145 
146     /// Internal helper for `build_decoder` to allocate an [`AvCodecContext`]. This needs to be
147     /// paired with a later call to [`AvCodecContext::init`].
alloc_context(&self) -> Result<AvCodecContext, AvCodecOpenError>148     fn alloc_context(&self) -> Result<AvCodecContext, AvCodecOpenError> {
149         // TODO(b:315859322): add safety doc string
150         #[allow(clippy::undocumented_unsafe_blocks)]
151         let context = unsafe { ffi::avcodec_alloc_context3(self.0).as_mut() }
152             .ok_or(AvCodecOpenError::ContextAllocation)?;
153 
154         Ok(AvCodecContext(context))
155     }
156 }
157 
158 /// A builder to create a [`AvCodecContext`] suitable for decoding.
159 // This struct wraps an AvCodecContext directly, but the only way it can be taken out is to call
160 // `build()`, which finalizes the context and prevent further modification to the callback, etc.
161 pub struct DecoderContextBuilder {
162     codec: *const ffi::AVCodec,
163     context: AvCodecContext,
164 }
165 
166 impl DecoderContextBuilder {
167     /// Set a custom callback that provides output buffers.
168     ///
169     /// `get_buffer2` is a function that decides which buffer is used to render a frame (see
170     /// libavcodec's documentation for `get_buffer2` for more details). If provided, this function
171     /// must be thread-safe.
172     /// `opaque` is a pointer that will be passed as first argument to `get_buffer2` when it is
173     /// called.
set_get_buffer_2( &mut self, get_buffer2: unsafe extern "C" fn(*mut ffi::AVCodecContext, *mut ffi::AVFrame, i32) -> i32, opaque: *mut libc::c_void, )174     pub fn set_get_buffer_2(
175         &mut self,
176         get_buffer2: unsafe extern "C" fn(*mut ffi::AVCodecContext, *mut ffi::AVFrame, i32) -> i32,
177         opaque: *mut libc::c_void,
178     ) {
179         // SAFETY:
180         // Safe because self.context.0 is a pointer to a live AVCodecContext allocation.
181         let context = unsafe { &mut *(self.context.0) };
182         context.get_buffer2 = Some(get_buffer2);
183         context.opaque = opaque;
184     }
185 
186     /// Build a decoder AvCodecContext from the configured options.
build(mut self) -> Result<AvCodecContext, AvCodecOpenError>187     pub fn build(mut self) -> Result<AvCodecContext, AvCodecOpenError> {
188         self.context.init(self.codec)?;
189         Ok(self.context)
190     }
191 }
192 
193 /// A builder to create a [`AvCodecContext`] suitable for encoding.
194 // This struct wraps an AvCodecContext directly, but the only way it can be taken out is to call
195 // `build()`, which finalizes the context and prevent further modification to the callback, etc.
196 pub struct EncoderContextBuilder {
197     codec: *const ffi::AVCodec,
198     context: AvCodecContext,
199 }
200 
201 impl EncoderContextBuilder {
202     /// Set the width of input frames for this encoding context.
set_dimensions(&mut self, dimensions: Dimensions)203     pub fn set_dimensions(&mut self, dimensions: Dimensions) {
204         // TODO(b:315859322): add safety doc string
205         #[allow(clippy::undocumented_unsafe_blocks)]
206         let context = unsafe { &mut *(self.context.0) };
207         context.width = dimensions.width as _;
208         context.height = dimensions.height as _;
209     }
210 
211     /// Set the time base for this encoding context.
set_time_base(&mut self, time_base: ffi::AVRational)212     pub fn set_time_base(&mut self, time_base: ffi::AVRational) {
213         // TODO(b:315859322): add safety doc string
214         #[allow(clippy::undocumented_unsafe_blocks)]
215         let context = unsafe { &mut *(self.context.0) };
216         context.time_base = time_base;
217     }
218 
219     /// Set the input pixel format for this encoding context.
set_pix_fmt(&mut self, fmt: AvPixelFormat)220     pub fn set_pix_fmt(&mut self, fmt: AvPixelFormat) {
221         // TODO(b:315859322): add safety doc string
222         #[allow(clippy::undocumented_unsafe_blocks)]
223         let context = unsafe { &mut *(self.context.0) };
224         context.pix_fmt = fmt.pix_fmt();
225     }
226 
227     /// Build a encoder AvCodecContext from the configured options.
build(mut self) -> Result<AvCodecContext, AvCodecOpenError>228     pub fn build(mut self) -> Result<AvCodecContext, AvCodecOpenError> {
229         self.context.init(self.codec)?;
230         Ok(self.context)
231     }
232 }
233 
234 impl Default for AvCodecIterator {
default() -> Self235     fn default() -> Self {
236         Self::new()
237     }
238 }
239 
240 /// Lightweight abstraction over libavcodec's `av_codec_iterate` function that can be used to
241 /// enumerate all the supported codecs.
242 pub struct AvCodecIterator(*mut libc::c_void);
243 
244 impl AvCodecIterator {
new() -> Self245     pub fn new() -> Self {
246         Self(std::ptr::null_mut())
247     }
248 }
249 
250 impl Iterator for AvCodecIterator {
251     type Item = AvCodec;
252 
next(&mut self) -> Option<Self::Item>253     fn next(&mut self) -> Option<Self::Item> {
254         // SAFETY:
255         // Safe because our pointer was initialized to `NULL` and we only use it with
256         // `av_codec_iterate`, which will update it to a valid value.
257         unsafe { ffi::av_codec_iterate(&mut self.0 as *mut *mut libc::c_void).as_ref() }
258             .map(AvCodec)
259     }
260 }
261 
262 /// Simple wrapper over `AVProfile` that provides helpful methods.
263 pub struct AvProfile(&'static ffi::AVProfile);
264 
265 impl AvProfile {
266     /// Return the profile id, which can be matched against FF_PROFILE_*.
profile(&self) -> u32267     pub fn profile(&self) -> u32 {
268         self.0.profile as u32
269     }
270 
271     /// Return the name of this profile.
name(&self) -> &'static str272     pub fn name(&self) -> &'static str {
273         const INVALID_PROFILE_STR: &str = "invalid profile";
274 
275         // SAFETY:
276         // Safe because `CStr::from_ptr` is called on a valid zero-terminated C string.
277         unsafe { CStr::from_ptr(self.0.name).to_str() }.unwrap_or(INVALID_PROFILE_STR)
278     }
279 }
280 
281 impl Display for AvProfile {
fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result282     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
283         f.write_str(self.name())
284     }
285 }
286 
287 impl Debug for AvProfile {
fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result288     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
289         Display::fmt(self, f)
290     }
291 }
292 
293 /// Lightweight abstraction over the array of supported profiles for a given codec.
294 pub struct AvProfileIterator(*const ffi::AVProfile);
295 
296 impl Iterator for AvProfileIterator {
297     type Item = AvProfile;
298 
next(&mut self) -> Option<Self::Item>299     fn next(&mut self) -> Option<Self::Item> {
300         // SAFETY:
301         // Safe because the contract of `new` stipulates we have received a valid `AVCodec`
302         // reference, thus the `profiles` pointer must either be NULL or point to a valid array
303         // or `VAProfile`s.
304         match unsafe { self.0.as_ref() } {
305             None => None,
306             Some(profile) => {
307                 match profile.profile {
308                     ffi::FF_PROFILE_UNKNOWN => None,
309                     _ => {
310                         // SAFETY:
311                         // Safe because we have been initialized to a static, valid profiles array
312                         // which is terminated by FF_PROFILE_UNKNOWN.
313                         self.0 = unsafe { self.0.offset(1) };
314                         Some(AvProfile(profile))
315                     }
316                 }
317             }
318         }
319     }
320 }
321 
322 #[derive(Clone, Copy)]
323 /// Simple wrapper over `AVPixelFormat` that provides helpful methods.
324 pub struct AvPixelFormat(ffi::AVPixelFormat);
325 
326 impl AvPixelFormat {
327     /// Return the name of this pixel format.
name(&self) -> &'static str328     pub fn name(&self) -> &'static str {
329         const INVALID_FORMAT_STR: &str = "invalid pixel format";
330 
331         // SAFETY:
332         // Safe because `av_get_pix_fmt_name` returns either NULL or a valid C string.
333         let pix_fmt_name = unsafe { ffi::av_get_pix_fmt_name(self.0) };
334         // SAFETY:
335         // Safe because `pix_fmt_name` is a valid pointer to a C string.
336         match unsafe {
337             pix_fmt_name
338                 .as_ref()
339                 .and_then(|s| CStr::from_ptr(s).to_str().ok())
340         } {
341             None => INVALID_FORMAT_STR,
342             Some(string) => string,
343         }
344     }
345 
346     /// Return the avcodec profile id, which can be matched against AV_PIX_FMT_*.
347     ///
348     /// Note that this is **not** the same as a fourcc.
pix_fmt(&self) -> ffi::AVPixelFormat349     pub fn pix_fmt(&self) -> ffi::AVPixelFormat {
350         self.0
351     }
352 
353     /// Return the fourcc of the pixel format, or a series of zeros if its fourcc is unknown.
fourcc(&self) -> [u8; 4]354     pub fn fourcc(&self) -> [u8; 4] {
355         // SAFETY:
356         // Safe because `avcodec_pix_fmt_to_codec_tag` does not take any pointer as input and
357         // handles any value passed as argument.
358         unsafe { ffi::avcodec_pix_fmt_to_codec_tag(self.0) }.to_le_bytes()
359     }
360 
361     /// Given the width and plane index, returns the line size (data pointer increment per row) in
362     /// bytes.
line_size(&self, width: u32, plane: usize) -> Result<usize, AvError>363     pub fn line_size(&self, width: u32, plane: usize) -> Result<usize, AvError> {
364         av_image_line_size(*self, width, plane)
365     }
366 
367     /// Given an iterator of line sizes and height, return the size required for each plane's buffer
368     /// in bytes.
plane_sizes<I: IntoIterator<Item = u32>>( &self, linesizes: I, height: u32, ) -> Result<Vec<usize>, AvError>369     pub fn plane_sizes<I: IntoIterator<Item = u32>>(
370         &self,
371         linesizes: I,
372         height: u32,
373     ) -> Result<Vec<usize>, AvError> {
374         av_image_plane_sizes(*self, linesizes, height)
375     }
376 }
377 
378 #[derive(Debug)]
379 pub struct FromAVPixelFormatError(());
380 
381 impl TryFrom<ffi::AVPixelFormat> for AvPixelFormat {
382     type Error = FromAVPixelFormatError;
383 
try_from(value: ffi::AVPixelFormat) -> Result<Self, Self::Error>384     fn try_from(value: ffi::AVPixelFormat) -> Result<Self, Self::Error> {
385         if value > ffi::AVPixelFormat_AV_PIX_FMT_NONE && value < ffi::AVPixelFormat_AV_PIX_FMT_NB {
386             Ok(AvPixelFormat(value))
387         } else {
388             Err(FromAVPixelFormatError(()))
389         }
390     }
391 }
392 
393 impl Display for AvPixelFormat {
fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result394     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
395         f.write_str(self.name())
396     }
397 }
398 
399 impl Debug for AvPixelFormat {
fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result400     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
401         let fourcc = self.fourcc();
402         f.write_fmt(format_args!(
403             "{}{}{}{}",
404             fourcc[0] as char, fourcc[1] as char, fourcc[2] as char, fourcc[3] as char
405         ))
406     }
407 }
408 
409 /// Lightweight abstraction over the array of supported pixel formats for a given codec.
410 pub struct AvPixelFormatIterator(*const ffi::AVPixelFormat);
411 
412 impl Iterator for AvPixelFormatIterator {
413     type Item = AvPixelFormat;
414 
next(&mut self) -> Option<Self::Item>415     fn next(&mut self) -> Option<Self::Item> {
416         // SAFETY:
417         // Safe because the contract of `AvCodec::new` and `AvCodec::pixel_format_iter` guarantees
418         // that we have been built from a valid `AVCodec` reference, which `pix_fmts` pointer
419         // must either be NULL or point to a valid array or `VAPixelFormat`s.
420         match unsafe { self.0.as_ref() } {
421             None => None,
422             Some(&pixfmt) => {
423                 match pixfmt {
424                     // Array of pixel formats is terminated by AV_PIX_FMT_NONE.
425                     ffi::AVPixelFormat_AV_PIX_FMT_NONE => None,
426                     _ => {
427                         // SAFETY:
428                         // Safe because we have been initialized to a static, valid profiles array
429                         // which is terminated by AV_PIX_FMT_NONE.
430                         self.0 = unsafe { self.0.offset(1) };
431                         Some(AvPixelFormat(pixfmt))
432                     }
433                 }
434             }
435         }
436     }
437 }
438 
439 /// A codec context from which decoding can be performed.
440 pub struct AvCodecContext(*mut ffi::AVCodecContext);
441 
442 impl Drop for AvCodecContext {
drop(&mut self)443     fn drop(&mut self) {
444         // SAFETY:
445         // Safe because our context member is properly allocated and owned by us.
446         // Note: `avcodec_open2` might not have been called in case we're wrapped by a
447         //       `DecoderContextBuilder` but avcodec_free_context works on both opened and closed
448         //       contexts.
449         unsafe { ffi::avcodec_free_context(&mut self.0) };
450     }
451 }
452 
453 impl AsRef<ffi::AVCodecContext> for AvCodecContext {
as_ref(&self) -> &ffi::AVCodecContext454     fn as_ref(&self) -> &ffi::AVCodecContext {
455         // SAFETY:
456         // Safe because our context member is properly initialized and fully owned by us.
457         unsafe { &*self.0 }
458     }
459 }
460 
461 pub enum TryReceiveResult {
462     Received,
463     TryAgain,
464     FlushCompleted,
465 }
466 
467 impl AvCodecContext {
468     /// Internal helper for [`DecoderContextBuilder`] to initialize the context.
init(&mut self, codec: *const ffi::AVCodec) -> Result<(), AvCodecOpenError>469     fn init(&mut self, codec: *const ffi::AVCodec) -> Result<(), AvCodecOpenError> {
470         // SAFETY:
471         // Safe because `codec` is a valid static AVCodec reference, and `self.0` is a valid
472         // AVCodecContext allocation.
473         if unsafe { ffi::avcodec_open2(self.0, codec, std::ptr::null_mut()) } < 0 {
474             return Err(AvCodecOpenError::ContextOpen);
475         }
476 
477         Ok(())
478     }
479 
480     /// Send a packet to be decoded by the codec.
481     ///
482     /// Returns `true` if the packet has been accepted and will be decoded, `false` if the codec can
483     /// not accept frames at the moment - in this case `try_receive_frame` must be called before
484     /// the packet can be submitted again.
485     ///
486     /// Error codes are the same as those returned by `avcodec_send_packet` with the exception of
487     /// EAGAIN which is converted into `Ok(false)` as it is not actually an error.
try_send_packet(&mut self, packet: &AvPacket) -> Result<bool, AvError>488     pub fn try_send_packet(&mut self, packet: &AvPacket) -> Result<bool, AvError> {
489         // SAFETY:
490         // Safe because the context is valid through the life of this object, and `packet`'s
491         // lifetime properties ensures its memory area is readable.
492         match unsafe { ffi::avcodec_send_packet(self.0, &packet.packet) } {
493             AVERROR_EAGAIN => Ok(false),
494             ret if ret >= 0 => Ok(true),
495             err => Err(AvError(err)),
496         }
497     }
498 
499     /// Attempt to write a decoded frame in `frame` if the codec has enough data to do so.
500     ///
501     /// Returned `Received` if `frame` has been filled with the next decoded frame, `TryAgain` if
502     /// no frame could be returned at that time (in which case `try_send_packet` should be called to
503     /// submit more input to decode), or `FlushCompleted` to signal that a previous flush triggered
504     /// by calling the `flush` method has completed.
505     ///
506     /// Error codes are the same as those returned by `avcodec_receive_frame` with the exception of
507     /// EAGAIN and EOF which are handled as `TryAgain` and `FlushCompleted` respectively.
try_receive_frame(&mut self, frame: &mut AvFrame) -> Result<TryReceiveResult, AvError>508     pub fn try_receive_frame(&mut self, frame: &mut AvFrame) -> Result<TryReceiveResult, AvError> {
509         // SAFETY:
510         // Safe because the context is valid through the life of this object, and `avframe` is
511         // guaranteed to contain a properly initialized frame.
512         match unsafe { ffi::avcodec_receive_frame(self.0, frame.0) } {
513             AVERROR_EAGAIN => Ok(TryReceiveResult::TryAgain),
514             AVERROR_EOF => Ok(TryReceiveResult::FlushCompleted),
515             ret if ret >= 0 => Ok(TryReceiveResult::Received),
516             err => Err(AvError(err)),
517         }
518     }
519 
520     /// Send a frame to be encoded by the codec.
521     ///
522     /// Returns `true` if the frame has been accepted and will be encoded, `false` if the codec can
523     /// not accept input at the moment - in this case `try_receive_frame` must be called before
524     /// the frame can be submitted again.
525     ///
526     /// Error codes are the same as those returned by `avcodec_send_frame` with the exception of
527     /// EAGAIN which is converted into `Ok(false)` as it is not actually an error.
try_send_frame(&mut self, frame: &AvFrame) -> Result<bool, AvError>528     pub fn try_send_frame(&mut self, frame: &AvFrame) -> Result<bool, AvError> {
529         // TODO(b:315859322): add safety doc string
530         #[allow(clippy::undocumented_unsafe_blocks)]
531         match unsafe { ffi::avcodec_send_frame(self.0, frame.0 as *const _) } {
532             AVERROR_EAGAIN => Ok(false),
533             ret if ret >= 0 => Ok(true),
534             err => Err(AvError(err)),
535         }
536     }
537 
538     /// Attempt to write an encoded frame in `packet` if the codec has enough data to do so.
539     ///
540     /// Returned `Received` if `packet` has been filled with encoded data, `TryAgain` if
541     /// no packet could be returned at that time (in which case `try_send_frame` should be called to
542     /// submit more input to decode), or `FlushCompleted` to signal that a previous flush triggered
543     /// by calling the `flush` method has completed.
544     ///
545     /// Error codes are the same as those returned by `avcodec_receive_packet` with the exception of
546     /// EAGAIN and EOF which are handled as `TryAgain` and `FlushCompleted` respectively.
try_receive_packet( &mut self, packet: &mut AvPacket, ) -> Result<TryReceiveResult, AvError>547     pub fn try_receive_packet(
548         &mut self,
549         packet: &mut AvPacket,
550     ) -> Result<TryReceiveResult, AvError> {
551         // SAFETY:
552         // Safe because the context is valid through the life of this object, and `avframe` is
553         // guaranteed to contain a properly initialized frame.
554         match unsafe { ffi::avcodec_receive_packet(self.0, &mut packet.packet) } {
555             AVERROR_EAGAIN => Ok(TryReceiveResult::TryAgain),
556             AVERROR_EOF => Ok(TryReceiveResult::FlushCompleted),
557             ret if ret >= 0 => Ok(TryReceiveResult::Received),
558             err => Err(AvError(err)),
559         }
560     }
561 
562     /// Reset the internal codec state/flush internal buffers.
563     /// Should be called e.g. when seeking or switching to a different stream.
reset(&mut self)564     pub fn reset(&mut self) {
565         // SAFETY:
566         // Safe because the context is valid through the life of this object.
567         unsafe { ffi::avcodec_flush_buffers(self.0) }
568     }
569 
570     /// Ask the context to start flushing, i.e. to process all pending input packets and produce
571     /// frames for them.
572     ///
573     /// The flush process is complete when `try_receive_frame` returns `FlushCompleted`,
flush_decoder(&mut self) -> Result<(), AvError>574     pub fn flush_decoder(&mut self) -> Result<(), AvError> {
575         // SAFETY:
576         // Safe because the context is valid through the life of this object.
577         AvError::result(unsafe { ffi::avcodec_send_packet(self.0, std::ptr::null()) })
578     }
579 
580     /// Ask the context to start flushing, i.e. to process all pending input frames and produce
581     /// packets for them.
582     ///
583     /// The flush process is complete when `try_receive_packet` returns `FlushCompleted`,
flush_encoder(&mut self) -> Result<(), AvError>584     pub fn flush_encoder(&mut self) -> Result<(), AvError> {
585         // SAFETY:
586         // Safe because the context is valid through the life of this object.
587         AvError::result(unsafe { ffi::avcodec_send_frame(self.0, std::ptr::null()) })
588     }
589 
590     /// Set the time base for this context.
set_time_base(&mut self, time_base: AVRational)591     pub fn set_time_base(&mut self, time_base: AVRational) {
592         // TODO(b:315859322): add safety doc string
593         #[allow(clippy::undocumented_unsafe_blocks)]
594         let context = unsafe { &mut *(self.0) };
595         context.time_base = time_base;
596     }
597 
598     /// Set the bit rate for this context.
set_bit_rate(&mut self, bit_rate: u64)599     pub fn set_bit_rate(&mut self, bit_rate: u64) {
600         // TODO(b:315859322): add safety doc string
601         #[allow(clippy::undocumented_unsafe_blocks)]
602         let context = unsafe { &mut *(self.0) };
603         context.bit_rate = bit_rate as _;
604     }
605 
606     /// Set the max bit rate (rc_max_rate) for this context.
set_max_bit_rate(&mut self, bit_rate: u64)607     pub fn set_max_bit_rate(&mut self, bit_rate: u64) {
608         // TODO(b:315859322): add safety doc string
609         #[allow(clippy::undocumented_unsafe_blocks)]
610         let context = unsafe { &mut *(self.0) };
611         context.rc_max_rate = bit_rate as _;
612     }
613 }
614 
615 /// Trait for types that can be used as data provider for a `AVBuffer`.
616 ///
617 /// `AVBuffer` is an owned buffer type, so all the type needs to do is being able to provide a
618 /// stable pointer to its own data as well as its length. Implementors need to be sendable across
619 /// threads because avcodec is allowed to use threads in its codec implementations.
620 pub trait AvBufferSource: Send {
as_ptr(&self) -> *const u8621     fn as_ptr(&self) -> *const u8;
as_mut_ptr(&mut self) -> *mut u8622     fn as_mut_ptr(&mut self) -> *mut u8 {
623         self.as_ptr() as *mut u8
624     }
len(&self) -> usize625     fn len(&self) -> usize;
is_empty(&self) -> bool626     fn is_empty(&self) -> bool;
627 }
628 
629 /// Wrapper around `AVBuffer` and `AVBufferRef`.
630 ///
631 /// libavcodec can manage its own memory for input and output data. Doing so implies a transparent
632 /// copy of user-provided data (packets or frames) from and to this memory, which is wasteful.
633 ///
634 /// This copy can be avoided by explicitly providing our own buffers to libavcodec using
635 /// `AVBufferRef`. Doing so means that the lifetime of these buffers becomes managed by avcodec.
636 /// This struct helps make this process safe by taking full ownership of an `AvBufferSource` and
637 /// dropping it when libavcodec is done with it.
638 pub struct AvBuffer(*mut ffi::AVBufferRef);
639 
640 impl AvBuffer {
641     /// Create a new `AvBuffer` from an `AvBufferSource`.
642     ///
643     /// Ownership of `source` is transferred to libavcodec, which will drop it when the number of
644     /// references to this buffer reaches zero.
645     ///
646     /// Returns `None` if the buffer could not be created due to an error in libavcodec.
new<D: AvBufferSource + 'static>(source: D) -> Option<Self>647     pub fn new<D: AvBufferSource + 'static>(source: D) -> Option<Self> {
648         // Move storage to the heap so we find it at the same place in `avbuffer_free`
649         let mut storage = Box::new(source);
650 
651         extern "C" fn avbuffer_free<D>(opaque: *mut c_void, _data: *mut u8) {
652             // SAFETY:
653             // Safe because `opaque` has been created from `Box::into_raw`. `storage` will be
654             // dropped immediately which will release any resources held by the storage.
655             let _ = unsafe { Box::from_raw(opaque as *mut D) };
656         }
657 
658         // SAFETY:
659         // Safe because storage points to valid data throughout the lifetime of AVBuffer and we are
660         // checking the return value against NULL, which signals an error.
661         Some(Self(unsafe {
662             ffi::av_buffer_create(
663                 storage.as_mut_ptr(),
664                 storage.len(),
665                 Some(avbuffer_free::<D>),
666                 Box::into_raw(storage) as *mut c_void,
667                 0,
668             )
669             .as_mut()?
670         }))
671     }
672 
673     /// Return a slice to the data contained in this buffer.
as_mut_slice(&mut self) -> &mut [u8]674     pub fn as_mut_slice(&mut self) -> &mut [u8] {
675         // SAFETY:
676         // Safe because the data has been initialized from valid storage in the constructor.
677         unsafe { std::slice::from_raw_parts_mut((*self.0).data, (*self.0).size) }
678     }
679 
680     /// Consumes the `AVBuffer`, returning a `AVBufferRef` that can be used in `AVFrame`, `AVPacket`
681     /// and others.
682     ///
683     /// After calling, the caller is responsible for unref-ing the returned AVBufferRef, either
684     /// directly or through one of the automatic management facilities in `AVFrame`, `AVPacket` or
685     /// others.
into_raw(self) -> *mut ffi::AVBufferRef686     pub fn into_raw(self) -> *mut ffi::AVBufferRef {
687         ManuallyDrop::new(self).0
688     }
689 }
690 
691 impl Drop for AvBuffer {
drop(&mut self)692     fn drop(&mut self) {
693         // SAFETY:
694         // Safe because `self.0` is a valid pointer to an AVBufferRef.
695         unsafe { ffi::av_buffer_unref(&mut self.0) };
696     }
697 }
698 
699 /// An encoded input packet that can be submitted to `AvCodecContext::try_send_packet`.
700 pub struct AvPacket<'a> {
701     packet: ffi::AVPacket,
702     _buffer_data: PhantomData<&'a ()>,
703 }
704 
705 impl<'a> Drop for AvPacket<'a> {
drop(&mut self)706     fn drop(&mut self) {
707         // SAFETY:
708         // Safe because `self.packet` is a valid `AVPacket` instance.
709         unsafe {
710             ffi::av_packet_unref(&mut self.packet);
711         }
712     }
713 }
714 
715 impl<'a> AsRef<ffi::AVPacket> for AvPacket<'a> {
as_ref(&self) -> &ffi::AVPacket716     fn as_ref(&self) -> &ffi::AVPacket {
717         &self.packet
718     }
719 }
720 
721 impl<'a> AvPacket<'a> {
722     /// Create an empty AvPacket without buffers.
723     ///
724     /// This packet should be only used with an encoder; in which case the encoder will
725     /// automatically allocate a buffer of appropriate size and store it inside this `AvPacket`.
empty() -> Self726     pub fn empty() -> Self {
727         Self {
728             packet: ffi::AVPacket {
729                 pts: AV_NOPTS_VALUE as i64,
730                 dts: AV_NOPTS_VALUE as i64,
731                 pos: -1,
732                 // SAFETY:
733                 // Safe because all the other elements of this struct can be zeroed.
734                 ..unsafe { std::mem::zeroed() }
735             },
736             _buffer_data: PhantomData,
737         }
738     }
739 
740     /// Create a new AvPacket that borrows the `input_data`.
741     ///
742     /// The returned `AvPacket` will hold a reference to `input_data`, meaning that libavcodec might
743     /// perform a copy from/to it.
new<T: AvBufferSource>(pts: i64, input_data: &'a mut T) -> Self744     pub fn new<T: AvBufferSource>(pts: i64, input_data: &'a mut T) -> Self {
745         Self {
746             packet: ffi::AVPacket {
747                 buf: std::ptr::null_mut(),
748                 pts,
749                 dts: AV_NOPTS_VALUE as i64,
750                 data: input_data.as_mut_ptr(),
751                 size: input_data.len() as c_int,
752                 side_data: std::ptr::null_mut(),
753                 pos: -1,
754                 // SAFETY:
755                 // Safe because all the other elements of this struct can be zeroed.
756                 ..unsafe { std::mem::zeroed() }
757             },
758             _buffer_data: PhantomData,
759         }
760     }
761 
762     /// Create a new AvPacket that owns the `av_buffer`.
763     ///
764     /// The returned `AvPacket` will have a `'static` lifetime and will keep `input_data` alive for
765     /// as long as libavcodec needs it.
new_owned(pts: i64, mut av_buffer: AvBuffer) -> Self766     pub fn new_owned(pts: i64, mut av_buffer: AvBuffer) -> Self {
767         let data_slice = av_buffer.as_mut_slice();
768         let data = data_slice.as_mut_ptr();
769         let size = data_slice.len() as i32;
770 
771         Self {
772             packet: ffi::AVPacket {
773                 buf: av_buffer.into_raw(),
774                 pts,
775                 dts: AV_NOPTS_VALUE as i64,
776                 data,
777                 size,
778                 side_data: std::ptr::null_mut(),
779                 pos: -1,
780                 // SAFETY:
781                 // Safe because all the other elements of this struct can be zeroed.
782                 ..unsafe { std::mem::zeroed() }
783             },
784             _buffer_data: PhantomData,
785         }
786     }
787 }
788 
789 /// An owned AVFrame, i.e. one decoded frame from libavcodec that can be converted into a
790 /// destination buffer.
791 pub struct AvFrame(*mut ffi::AVFrame);
792 
793 /// A builder for AVFrame that allows specifying buffers and image metadata.
794 pub struct AvFrameBuilder(AvFrame);
795 
796 /// A descriptor describing a subslice of `buffers` in [`AvFrameBuilder::build_owned`] that
797 /// represents a plane's image data.
798 pub struct PlaneDescriptor {
799     /// The index within `buffers`.
800     pub buffer_index: usize,
801     /// The offset from the start of `buffers[buffer_index]`.
802     pub offset: usize,
803     /// The increment of data pointer in bytes per row of the plane.
804     pub stride: usize,
805 }
806 
807 #[derive(Debug, ThisError)]
808 pub enum AvFrameError {
809     #[error("failed to allocate AVFrame object")]
810     FrameAllocationFailed,
811     #[error("dimension is negative or too large")]
812     DimensionOverflow,
813     #[error("a row does not fit in the specified stride")]
814     InvalidStride,
815     #[error("buffer index out of range")]
816     BufferOutOfRange,
817     #[error("specified dimensions overflow the buffer size")]
818     BufferTooSmall,
819     #[error("plane reference to buffer alias each other")]
820     BufferAlias,
821     #[error("error while calling libavcodec")]
822     AvError(#[from] AvError),
823 }
824 
825 impl AvFrame {
826     /// Create a new AvFrame. The frame's parameters and backing memory will be assigned when it is
827     /// decoded into.
new() -> Result<Self, AvFrameError>828     pub fn new() -> Result<Self, AvFrameError> {
829         Ok(Self(
830             // SAFETY:
831             // Safe because `av_frame_alloc` does not take any input.
832             unsafe { ffi::av_frame_alloc().as_mut() }.ok_or(AvFrameError::FrameAllocationFailed)?,
833         ))
834     }
835 
836     /// Create a new AvFrame builder that allows setting the frame's parameters and backing memory
837     /// through its methods.
builder() -> Result<AvFrameBuilder, AvFrameError>838     pub fn builder() -> Result<AvFrameBuilder, AvFrameError> {
839         AvFrame::new().map(AvFrameBuilder)
840     }
841 
842     /// Return the frame's width and height.
dimensions(&self) -> Dimensions843     pub fn dimensions(&self) -> Dimensions {
844         Dimensions {
845             width: self.as_ref().width as _,
846             height: self.as_ref().height as _,
847         }
848     }
849 
850     /// Return the frame's pixel format.
format(&self) -> AvPixelFormat851     pub fn format(&self) -> AvPixelFormat {
852         AvPixelFormat(self.as_ref().format)
853     }
854 
855     /// Set the picture type (I-frame, P-frame etc.) on this frame.
set_pict_type(&mut self, ty: AVPictureType)856     pub fn set_pict_type(&mut self, ty: AVPictureType) {
857         // SAFETY:
858         // Safe because self.0 is a valid AVFrame reference.
859         unsafe {
860             (*self.0).pict_type = ty;
861         }
862     }
863 
864     /// Set the presentation timestamp (PTS) of this frame.
set_pts(&mut self, ts: i64)865     pub fn set_pts(&mut self, ts: i64) {
866         // SAFETY:
867         // Safe because self.0 is a valid AVFrame reference.
868         unsafe {
869             (*self.0).pts = ts;
870         }
871     }
872 
873     /// Query if this AvFrame is writable, i.e. it is refcounted and the refcounts are 1.
is_writable(&self) -> bool874     pub fn is_writable(&self) -> bool {
875         // SAFETY:
876         // Safe because self.0 is a valid AVFrame reference.
877         unsafe { ffi::av_frame_is_writable(self.0) != 0 }
878     }
879 
880     /// If the frame is not writable already (see [`is_writable`]), make a copy of its buffer to
881     /// make it writable.
882     ///
883     /// [`is_writable`]: AvFrame::is_writable
make_writable(&mut self) -> Result<(), AvFrameError>884     pub fn make_writable(&mut self) -> Result<(), AvFrameError> {
885         // SAFETY:
886         // Safe because self.0 is a valid AVFrame reference.
887         AvError::result(unsafe { ffi::av_frame_make_writable(self.0) }).map_err(Into::into)
888     }
889 }
890 
891 impl AvFrameBuilder {
892     /// Set the frame's width and height.
893     ///
894     /// The dimensions must not be greater than `i32::MAX`.
set_dimensions(&mut self, dimensions: Dimensions) -> Result<(), AvFrameError>895     pub fn set_dimensions(&mut self, dimensions: Dimensions) -> Result<(), AvFrameError> {
896         // SAFETY:
897         // Safe because self.0 is a valid AVFrame instance and width and height are in range.
898         unsafe {
899             (*self.0 .0).width = dimensions
900                 .width
901                 .try_into()
902                 .map_err(|_| AvFrameError::DimensionOverflow)?;
903             (*self.0 .0).height = dimensions
904                 .height
905                 .try_into()
906                 .map_err(|_| AvFrameError::DimensionOverflow)?;
907         }
908         Ok(())
909     }
910 
911     /// Set the frame's format.
set_format(&mut self, format: AvPixelFormat) -> Result<(), AvFrameError>912     pub fn set_format(&mut self, format: AvPixelFormat) -> Result<(), AvFrameError> {
913         // SAFETY:
914         // Safe because self.0 is a valid AVFrame instance and format is a valid pixel format.
915         unsafe {
916             (*self.0 .0).format = format.pix_fmt();
917         }
918         Ok(())
919     }
920 
921     /// Build an AvFrame from iterators of [`AvBuffer`]s and subslice of buffers describing the
922     /// planes.
923     ///
924     /// The frame will own the `buffers`.
925     ///
926     /// This function checks that:
927     /// - Each plane fits inside the bounds of the associated buffer.
928     /// - Different planes do not overlap each other's buffer slice. In this check, all planes are
929     ///   assumed to be potentially mutable, regardless of whether the AvFrame is actually used for
930     ///   read or write access. Aliasing reference to the same buffer will be rejected, since it can
931     ///   potentially allow routines to overwrite each
932     //    other's result.
933     ///   An exception to this is when the same buffer is passed multiple times in `buffers`. In
934     ///   this case, each buffer is treated as a different buffer. Since clones have to be made to
935     ///   be passed multiple times in `buffers`, the frame will not be considered [writable]. Hence
936     ///   aliasing is safe in this case, but the caller is required to explicit opt-in to this
937     ///   read-only handling by passing clones of the buffer into `buffers` and have a different
938     ///   buffer index for each plane combination that could overlap in their range.
939     ///
940     /// [writable]: AvFrame::is_writable
build_owned< BI: IntoIterator<Item = AvBuffer>, PI: IntoIterator<Item = PlaneDescriptor>, >( self, buffers: BI, planes: PI, ) -> Result<AvFrame, AvFrameError>941     pub fn build_owned<
942         BI: IntoIterator<Item = AvBuffer>,
943         PI: IntoIterator<Item = PlaneDescriptor>,
944     >(
945         self,
946         buffers: BI,
947         planes: PI,
948     ) -> Result<AvFrame, AvFrameError> {
949         let mut buffers: Vec<_> = buffers.into_iter().collect();
950         let planes: Vec<_> = planes.into_iter().collect();
951         let format = self.0.format();
952         let plane_sizes = format.plane_sizes(
953             planes.iter().map(|x| x.stride as u32),
954             self.0.dimensions().height,
955         )?;
956         let mut ranges = vec![];
957 
958         for (
959             plane,
960             PlaneDescriptor {
961                 buffer_index,
962                 offset,
963                 stride,
964             },
965         ) in planes.into_iter().enumerate()
966         {
967             if buffer_index > buffers.len() {
968                 return Err(AvFrameError::BufferOutOfRange);
969             }
970             let end = offset + plane_sizes[plane];
971             if end > buffers[buffer_index].as_mut_slice().len() {
972                 return Err(AvFrameError::BufferTooSmall);
973             }
974             if stride < format.line_size(self.0.dimensions().width, plane)? {
975                 return Err(AvFrameError::InvalidStride);
976             }
977             // TODO(b:315859322): add safety doc string
978             #[allow(clippy::undocumented_unsafe_blocks)]
979             unsafe {
980                 (*self.0 .0).data[plane] =
981                     buffers[buffer_index].as_mut_slice()[offset..].as_mut_ptr();
982                 (*self.0 .0).linesize[plane] = stride as c_int;
983             }
984             ranges.push((buffer_index, offset, end));
985         }
986 
987         // Check for range overlaps.
988         // See function documentation for the exact rule and reasoning.
989         ranges.sort_unstable();
990         for pair in ranges.windows(2) {
991             // (buffer_index, start, end)
992             let (b0, _s0, e0) = pair[0];
993             let (b1, s1, _e1) = pair[1];
994 
995             if b0 != b1 {
996                 continue;
997             }
998             // Note that s0 <= s1 always holds, so we only need to check
999             // that the start of the second range is before the end of the first range.
1000             if s1 < e0 {
1001                 return Err(AvFrameError::BufferAlias);
1002             }
1003         }
1004 
1005         for (i, buf) in buffers.into_iter().enumerate() {
1006             // SAFETY:
1007             // Safe because self.0 is a valid AVFrame instance and buffers contains valid AvBuffers.
1008             unsafe {
1009                 (*self.0 .0).buf[i] = buf.into_raw();
1010             }
1011         }
1012         Ok(self.0)
1013     }
1014 }
1015 
1016 impl AsRef<ffi::AVFrame> for AvFrame {
as_ref(&self) -> &ffi::AVFrame1017     fn as_ref(&self) -> &ffi::AVFrame {
1018         // SAFETY:
1019         // Safe because the AVFrame has been properly initialized during construction.
1020         unsafe { &*self.0 }
1021     }
1022 }
1023 
1024 impl Deref for AvFrame {
1025     type Target = ffi::AVFrame;
1026 
deref(&self) -> &Self::Target1027     fn deref(&self) -> &Self::Target {
1028         // SAFETY:
1029         // Safe because the AVFrame has been properly initialized during construction.
1030         unsafe { self.0.as_ref().unwrap() }
1031     }
1032 }
1033 
1034 impl Drop for AvFrame {
drop(&mut self)1035     fn drop(&mut self) {
1036         // SAFETY:
1037         // Safe because the AVFrame is valid through the life of this object and fully owned by us.
1038         unsafe { ffi::av_frame_free(&mut self.0) };
1039     }
1040 }
1041 
1042 #[cfg(test)]
1043 mod tests {
1044     use std::sync::atomic::AtomicBool;
1045     use std::sync::atomic::Ordering;
1046     use std::sync::Arc;
1047 
1048     use super::*;
1049 
1050     #[test]
test_averror()1051     fn test_averror() {
1052         // Just test that the error is wrapper properly. The bindings test module already checks
1053         // that the error bindings correspond to the right ffmpeg errors.
1054         let averror = AvError(AVERROR_EOF);
1055         let msg = format!("{}", averror);
1056         assert_eq!(msg, "End of file");
1057 
1058         let averror = AvError(0);
1059         let msg = format!("{}", averror);
1060         assert_eq!(msg, "Success");
1061 
1062         let averror = AvError(10);
1063         let msg = format!("{}", averror);
1064         assert_eq!(msg, "Unknown avcodec error 10");
1065     }
1066 
1067     // Test that the AVPacket wrapper frees the owned AVBuffer on drop.
1068     #[test]
test_avpacket_drop()1069     fn test_avpacket_drop() {
1070         struct DropTestBufferSource {
1071             dropped: Arc<AtomicBool>,
1072         }
1073         impl Drop for DropTestBufferSource {
1074             fn drop(&mut self) {
1075                 self.dropped.store(true, Ordering::SeqCst);
1076             }
1077         }
1078         impl AvBufferSource for DropTestBufferSource {
1079             fn as_ptr(&self) -> *const u8 {
1080                 [].as_ptr()
1081             }
1082 
1083             fn len(&self) -> usize {
1084                 0
1085             }
1086 
1087             fn is_empty(&self) -> bool {
1088                 true
1089             }
1090         }
1091 
1092         let dropped = Arc::new(AtomicBool::new(false));
1093 
1094         let pkt = AvPacket::new_owned(
1095             0,
1096             AvBuffer::new(DropTestBufferSource {
1097                 dropped: dropped.clone(),
1098             })
1099             .unwrap(),
1100         );
1101         assert!(!dropped.load(Ordering::SeqCst));
1102         drop(pkt);
1103         assert!(dropped.load(Ordering::SeqCst));
1104     }
1105 }
1106