Make the decoder more robust against extra trailing bits in the NALUs

src/nal/sps.rs

@@ -2116,6 +2116,60 @@ mod test {
         1080,
         25.0; "1920x1080 hikvision nal hrd + vcl hrd"
     )]
+    #[test_case(
+        vec![
+            103, 100, 0, 21, 172, 228, 5, 132,
+            183, 254, 1, 0, 0, 234, 32, 0,
+            0, 3, 0, 32, 0, 0, 7, 147,
+            226, 133, 73, 0, 10
+        ],
+        SeqParameterSet{
+            profile_idc: ProfileIdc::from(100),
+            constraint_flags: ConstraintFlags::from(0),
+            level_idc: 21,
+            seq_parameter_set_id: SeqParamSetId::from_u32(0).unwrap(),
+            chroma_info: ChromaInfo{
+                chroma_format: ChromaFormat::YUV420,
+                ..ChromaInfo::default()
+            },
+            log2_max_frame_num_minus4: 0,
+            pic_order_cnt: PicOrderCntType::TypeZero {
+                log2_max_pic_order_cnt_lsb_minus4: 0
+            },
+            max_num_ref_frames: 3,
+            gaps_in_frame_num_value_allowed_flag: false,
+            pic_width_in_mbs_minus1: 21,
+            pic_height_in_map_units_minus1: 8,
+            frame_mbs_flags: FrameMbsFlags::Fields {
+                mb_adaptive_frame_field_flag: true
+            },
+            direct_8x8_inference_flag: true,
+            frame_cropping: None,
+            vui_parameters: Some(VuiParameters{
+                aspect_ratio_info: Some(AspectRatioInfo::Extended(128, 117)),
+                video_signal_type: None,
+                timing_info: Some(TimingInfo{
+                    num_units_in_tick: 1,
+                    time_scale: 60,
+                    fixed_frame_rate_flag: true,
+                }),
+                pic_struct_present_flag: true,
+                bitstream_restrictions: Some(BitstreamRestrictions {
+                    motion_vectors_over_pic_boundaries_flag: true,
+                    max_bytes_per_pic_denom: 0,
+                    max_bits_per_mb_denom: 0,
+                    log2_max_mv_length_horizontal: 9,
+                    log2_max_mv_length_vertical: 9,
+                    max_num_reorder_frames: 0,
+                    max_dec_frame_buffering: 3
+                }),
+                ..VuiParameters::default()
+            }),
+        },
+        352,
+        288,
+        30.0; "crew_cif (with extra trailing bytes)"
+    )]
     fn test_sps(byts: Vec<u8>, sps: SeqParameterSet, width: u32, height: u32, fps: f64) {
         let sps_rbsp = decode_nal(&byts).unwrap();
         let sps2 = SeqParameterSet::from_bits(BitReader::new(&*sps_rbsp)).unwrap();

src/rbsp.rs

@@ -23,6 +23,7 @@
 //! the sequence `0x00 0x00 0x03` with `0x00 0x00`).
 
 use bitstream_io::read::BitRead as _;
+use log::warn;
 use std::borrow::Cow;
 use std::io::BufRead;
 use std::io::Read;
@@ -260,8 +261,10 @@ pub enum BitReaderError {
     /// An Exp-Golomb-coded syntax elements value has more than 32 bits.
     ExpGolombTooLarge(&'static str),
 
-    /// The stream was positioned before the final one bit on [BitRead::finish_rbsp].
-    RemainingData,
+    /// The stream did not have a stop bit equal to one when finishing a stream.
+    StopBitNotOne,
+    /// The stream did not have all alignment bits equal to zero when finishing a stream.
+    AlignmentBitNotZero,
 
     Unaligned,
 }
@@ -397,37 +400,81 @@ impl<R: std::io::BufRead + Clone> BitRead for BitReader<R> {
     }
 
     fn finish_rbsp(mut self) -> Result<(), BitReaderError> {
-        // The next bit is expected to be the final one bit.
+        // rbsp_trailing_bits( ) {
+        //     rbsp_stop_one_bit /* equal to 1 */
+        //     while( !byte_aligned( ) )
+        //         rbsp_alignment_zero_bit /* equal to 0 */
+        // }
+
+        // consume `rbsp_stop_one_bit`
         if !self
             .reader
             .read_bit()
-            .map_err(|e| BitReaderError::ReaderErrorFor("finish", e))?
+            .map_err(|e| BitReaderError::ReaderErrorFor("rbsp_stop_one_bit", e))?
         {
-            // It was a zero! Determine if we're past the end or haven't reached it yet.
-            match self.reader.read_unary1() {
-                Err(e) => return Err(BitReaderError::ReaderErrorFor("finish", e)),
-                Ok(_) => return Err(BitReaderError::RemainingData),
+            return Err(BitReaderError::StopBitNotOne);
+        }
+
+        // consume `rbsp_alignment_zero_bit`(s)
+        while !self.reader.byte_aligned() {
+            match self.reader.read_bit() {
+                Ok(true) => return Err(BitReaderError::AlignmentBitNotZero),
+                Ok(false) => continue,
+                Err(e) => return Err(BitReaderError::ReaderErrorFor("rbsp_alignment_zero_bit", e)),
             }
         }
-        // All remaining bits in the stream must then be zeros.
-        match self.reader.read_unary1() {
+
+        // verify that there's nothing past the end `rbsp_trailing_bits`
+        match self.reader.read_bit() {
+            Ok(_) => {
+                // this branch only gets taken if there is extra data after `rbsp_trailing_bits`
+                warn!("BitReader: extra data after `rbsp_trailing_bits`");
+                Ok(())
+            }
             Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => Ok(()),
             Err(e) => Err(BitReaderError::ReaderErrorFor("finish", e)),
-            Ok(_) => Err(BitReaderError::RemainingData),
         }
     }
 
     fn finish_sei_payload(mut self) -> Result<(), BitReaderError> {
-        match self.reader.read_bit() {
-            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(()),
-            Err(e) => return Err(BitReaderError::ReaderErrorFor("finish", e)),
-            Ok(false) => return Err(BitReaderError::RemainingData),
-            Ok(true) => {}
+        // if( !byte_aligned( ) ) {
+        //     bit_equal_to_one /* equal to 1 */
+        //     while( !byte_aligned( ) )
+        //         bit_equal_to_zero /* equal to 0 */
+        // }
+
+        if !self.reader.byte_aligned() {
+            // consume `bit_equal_to_one`
+            if !self
+                .reader
+                .read_bit()
+                .map_err(|e| BitReaderError::ReaderErrorFor("bit_equal_to_one", e))?
+            {
+                return Err(BitReaderError::StopBitNotOne);
+            }
+            while !self.reader.byte_aligned() {
+                // consume `bit_equal_to_zero`(s)
+                while !self.reader.byte_aligned() {
+                    match self.reader.read_bit() {
+                        Ok(true) => return Err(BitReaderError::AlignmentBitNotZero),
+                        Ok(false) => continue,
+                        Err(e) => {
+                            return Err(BitReaderError::ReaderErrorFor("bit_equal_to_zero", e))
+                        }
+                    }
+                }
+            }
         }
-        match self.reader.read_unary1() {
+
+        // verify that there's nothing past the end
+        match self.reader.read_bit() {
+            Ok(_) => {
+                // this branch only gets taken if there is extra data after the SEI payload
+                warn!("BitReader: extra data after SEI payload");
+                Ok(())
+            }
             Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => Ok(()),
             Err(e) => Err(BitReaderError::ReaderErrorFor("finish", e)),
-            Ok(_) => Err(BitReaderError::RemainingData),
         }
     }
 }