diff --git a/crates/store/re_video/src/decode/ffmpeg_h264/sps.rs b/crates/store/re_video/src/decode/ffmpeg_h264/sps.rs index b14163c8cc02..2500a1494839 100644 --- a/crates/store/re_video/src/decode/ffmpeg_h264/sps.rs +++ b/crates/store/re_video/src/decode/ffmpeg_h264/sps.rs @@ -2,7 +2,7 @@ use crate::decode::YuvPixelLayout; use super::nalu::{NalHeader, NalUnitType}; -#[derive(thiserror::Error, Debug)] +#[derive(thiserror::Error, Debug, PartialEq, Eq)] pub enum SpsParsingError { #[error("SPS buffer too small")] UnexpectedEndOfSpsBuffer, @@ -271,16 +271,13 @@ fn read_bits( ) -> Result { debug_assert!(num_bits <= 32); - let byte_pos = *bit_read_pos / 8; - if buffer.len() <= byte_pos { - return Err(SpsParsingError::UnexpectedEndOfSpsBuffer); - } - - let num_bytes = num_bits.next_multiple_of(8) / 8; - if buffer.len() < byte_pos + num_bytes { + let highest_byte_read = (num_bits + *bit_read_pos).next_multiple_of(8) / 8; + if buffer.len() < highest_byte_read { return Err(SpsParsingError::UnexpectedEndOfSpsBuffer); } + // Read bit by bit. + // Obviously this can be sped up by reading bytes when possible, but let's keep it simple. let mut result = 0; for n in 0..num_bits { let bit_pos = *bit_read_pos + n; @@ -318,6 +315,8 @@ fn read_exponential_golomb( #[cfg(test)] mod tests { + use crate::decode::ffmpeg_h264::sps::SpsParsingError; + use super::{read_bits, read_exponential_golomb}; #[test] @@ -337,7 +336,14 @@ mod tests { ); assert_eq!(bit_pos, 12); - assert_eq!(read_bits(&mut 0, &[0], 9), Err(())); + assert_eq!( + read_bits(&mut 0, &[0], 9), + Err(SpsParsingError::UnexpectedEndOfSpsBuffer) + ); + assert_eq!( + read_bits(&mut 1, &[0], 8), + Err(SpsParsingError::UnexpectedEndOfSpsBuffer) + ); } #[test] @@ -369,5 +375,10 @@ mod tests { 6 ); assert_eq!(bit_pos, 7); + + assert_eq!( + read_exponential_golomb(&mut 0, &[0b_0000_1111]), + Err(SpsParsingError::UnexpectedEndOfSpsBuffer) + ); } }