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