From 549d4e6dc8e06bcf16b692ce8b4d3c86749132c8 Mon Sep 17 00:00:00 2001 From: David Holroyd Date: Mon, 26 Feb 2024 00:33:34 +0000 Subject: [PATCH] Bounds checks for bistream_restriction fields Fixes #55 --- Cargo.toml | 1 + src/lib.rs | 40 +++++ src/nal/sps.rs | 418 +++++++++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 442 insertions(+), 17 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 5c68651..5e1a331 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,7 @@ hex-slice = "0.1.4" memchr = "2.1.1" rfc6381-codec = "0.1" log = "0.4" +lazy_static = "1.4" [dev-dependencies] hex-literal = "0.4.1" diff --git a/src/lib.rs b/src/lib.rs index 2755568..8991ade 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ #![forbid(unsafe_code)] #![deny(rust_2018_idioms)] +use crate::rbsp::BitReaderError; use std::fmt::Debug; pub mod annexb; @@ -85,6 +86,45 @@ impl Debug for ParamSetMap { } } +#[derive(Debug)] +pub enum FieldErr { + TooLarge { value: u32, max_value: u32 }, + TooSmall { value: u32, min_value: u32 }, + BitReader(BitReaderError), +} + +trait U32Field { + fn max(self, max_value: u32) -> Result; + fn min(self, min_value: u32) -> Result; +} +impl U32Field for Result { + fn max(self, max_value: u32) -> Result { + match self { + Ok(value) => { + if value > max_value { + Err(FieldErr::TooLarge { value, max_value }) + } else { + Ok(value) + } + } + Err(e) => Err(FieldErr::BitReader(e)), + } + } + + fn min(self, min_value: u32) -> Result { + match self { + Ok(value) => { + if value < min_value { + Err(FieldErr::TooSmall { value, min_value }) + } else { + Ok(value) + } + } + Err(e) => Err(FieldErr::BitReader(e)), + } + } +} + #[cfg(test)] mod tests { #[test] diff --git a/src/nal/sps.rs b/src/nal/sps.rs index b67c9d5..19aef52 100644 --- a/src/nal/sps.rs +++ b/src/nal/sps.rs @@ -1,8 +1,12 @@ +use crate::U32Field; use crate::{ nal::pps::{ParamSetId, ParamSetIdError}, rbsp::{BitRead, BitReaderError}, + FieldErr, }; +use lazy_static::lazy_static; use std::fmt::{self, Debug}; +use std::num::NonZeroU8; #[derive(Debug)] pub enum SpsError { @@ -19,6 +23,7 @@ pub enum SpsError { name: &'static str, value: u32, }, + Field(FieldErr), /// The frame-cropping values are too large vs. the coded picture size, CroppingError(FrameCropping), /// The `cpb_cnt_minus1` field must be between 0 and 31 inclusive. @@ -141,7 +146,7 @@ impl Debug for ConstraintFlags { } } -#[derive(Debug, PartialEq)] +#[derive(Debug, PartialEq, Hash, Eq)] #[allow(non_camel_case_types)] pub enum Level { Unknown(u8), @@ -162,6 +167,9 @@ pub enum Level { L5, L5_1, L5_2, + L6, + L6_1, + L6_2, } impl Level { pub fn from_constraint_flags_and_level_idc( @@ -191,6 +199,9 @@ impl Level { 50 => Level::L5, 51 => Level::L5_1, 52 => Level::L5_2, + 60 => Level::L6, + 61 => Level::L6_1, + 62 => Level::L6_2, _ => Level::Unknown(level_idc), } } @@ -212,6 +223,9 @@ impl Level { Level::L5 => 50, Level::L5_1 => 51, Level::L5_2 => 52, + Level::L6 => 60, + Level::L6_1 => 61, + Level::L6_2 => 62, Level::Unknown(level_idc) => level_idc, } } @@ -797,18 +811,49 @@ pub struct BitstreamRestrictions { pub max_dec_frame_buffering: u32, } impl BitstreamRestrictions { - fn read(r: &mut R) -> Result, BitReaderError> { - let bitstream_restriction_flag = r.read_bool("bitstream_restriction_flag")?; + fn read( + r: &mut R, + sps: &SeqParameterSet, + ) -> Result, FieldErr> { + let bitstream_restriction_flag = r + .read_bool("bitstream_restriction_flag") + .map_err(FieldErr::BitReader)?; Ok(if bitstream_restriction_flag { + let motion_vectors_over_pic_boundaries_flag = r + .read_bool("motion_vectors_over_pic_boundaries_flag") + .map_err(FieldErr::BitReader)?; + let max_bytes_per_pic_denom = r.read_ue("max_bytes_per_pic_denom").max(16)?; + let max_bits_per_mb_denom = r.read_ue("max_bits_per_mb_denom").max(16)?; + let log2_max_mv_length_horizontal = + r.read_ue("log2_max_mv_length_horizontal").max(15)?; + let log2_max_mv_length_vertical = r.read_ue("log2_max_mv_length_vertical").max(15)?; + let max_num_reorder_frames = r + .read_ue("max_num_reorder_frames") + .map_err(FieldErr::BitReader)?; + let max_dec_frame_buffering = r + .read_ue("max_dec_frame_buffering") + .min(max_num_reorder_frames)?; + if max_num_reorder_frames > max_dec_frame_buffering { + return Err(FieldErr::TooLarge { + value: max_num_reorder_frames, + max_value: max_dec_frame_buffering, + }); + } + let max = max_val_for_max_dec_frame_buffering(sps); + if max_dec_frame_buffering > max { + return Err(FieldErr::TooLarge { + value: max_dec_frame_buffering, + max_value: max, + }); + } Some(BitstreamRestrictions { - motion_vectors_over_pic_boundaries_flag: r - .read_bool("motion_vectors_over_pic_boundaries_flag")?, - max_bytes_per_pic_denom: r.read_ue("max_bytes_per_pic_denom")?, - max_bits_per_mb_denom: r.read_ue("max_bits_per_mb_denom")?, - log2_max_mv_length_horizontal: r.read_ue("log2_max_mv_length_horizontal")?, - log2_max_mv_length_vertical: r.read_ue("log2_max_mv_length_vertical")?, - max_num_reorder_frames: r.read_ue("max_num_reorder_frames")?, - max_dec_frame_buffering: r.read_ue("max_dec_frame_buffering")?, + motion_vectors_over_pic_boundaries_flag, + max_bytes_per_pic_denom, + max_bits_per_mb_denom, + log2_max_mv_length_horizontal, + log2_max_mv_length_vertical, + max_num_reorder_frames, + max_dec_frame_buffering, }) } else { None @@ -816,6 +861,55 @@ impl BitstreamRestrictions { } } +// calculates the maximum allowed value for the max_dec_frame_buffering field +fn max_val_for_max_dec_frame_buffering(sps: &SeqParameterSet) -> u32 { + let level = Level::from_constraint_flags_and_level_idc( + ConstraintFlags::from(sps.constraint_flags), + sps.level_idc, + ); + let profile = Profile::from_profile_idc(sps.profile_idc); + let pic_width_in_mbs = sps.pic_width_in_mbs_minus1 + 1; + let pic_height_in_map_units = sps.pic_height_in_map_units_minus1 + 1; + let frame_height_in_mbs = if let FrameMbsFlags::Frames = sps.frame_mbs_flags { + 1 + } else { + 2 + } * pic_height_in_map_units; + let max_dpb_mbs = LEVEL_LIMITS.get(&level).unwrap().max_dpb_mbs; + match profile { + // "A.3.1 - Level limits common to the Baseline, Constrained Baseline, Main, and Extended + // profiles" + Profile::Baseline | Profile::Main | Profile::Extended => { + std::cmp::min(max_dpb_mbs / (pic_width_in_mbs * frame_height_in_mbs), 16) + } + // "A.3.2 - Level limits common to the High, Progressive High, Constrained High, High 10, + // Progressive High 10, High 4:2:2, High 4:4:4 Predictive, High 10 Intra, High 4:2:2 Intra, + // High 4:4:4 Intra, and CAVLC 4:4:4 Intra profiles" + Profile::High | Profile::High422 | Profile::High10 | Profile::High444 => { + std::cmp::min(max_dpb_mbs / (pic_width_in_mbs * frame_height_in_mbs), 16) + } + + // "G.10.2.1 - Level limits common to Scalable Baseline, Scalable Constrained Baseline, + // Scalable High, Scalable Constrained High, and Scalable High Intra profiles" + Profile::ScalableBase | Profile::ScalableHigh => { + // Min( MaxDpbMbs / ( PicWidthInMbs * FrameHeightInMbs ), 16 ) + std::cmp::min(max_dpb_mbs / (pic_width_in_mbs * frame_height_in_mbs), 16) + } + + // "H.10.2.1 - Level limits common to Multiview High, Stereo High, and MFC High profiles" + //Profile::MultiviewHigh | Profile::StereoHigh | Profile::MFCDepthHigh => { + // // Min( mvcScaleFactor * MaxDpbMbs / ( PicWidthInMbs * FrameHeightInMbs ), Max( 1, Ceil( log2( NumViews ) ) ) * 16 ) + //} + + // "I.10.2.1 - Level limits common to Multiview Depth High profiles" + //Profile::MultiviewDepthHigh | Profile::EnhancedMultiviewDepthHigh => { + // let mvcd_scale_factor = 2.5; + // std::cmp::min( mvcd_scale_factor * max_dpb_mbs / ( TotalPicSizeInMbs / NumViews ) ), std::cmp::max(1, Ceil( log2( NumViews ) ) ) * 16 ) + //} + _ => unimplemented!("{:?}", profile), + } +} + #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct VuiParameters { pub aspect_ratio_info: Option, @@ -830,7 +924,10 @@ pub struct VuiParameters { pub bitstream_restrictions: Option, } impl VuiParameters { - fn read(r: &mut R) -> Result, SpsError> { + fn read( + r: &mut R, + sps: &SeqParameterSet, + ) -> Result, SpsError> { let vui_parameters_present_flag = r.read_bool("vui_parameters_present_flag")?; Ok(if vui_parameters_present_flag { let mut hrd_parameters_present = false; @@ -848,7 +945,8 @@ impl VuiParameters { None }, pic_struct_present_flag: r.read_bool("pic_struct_present_flag")?, - bitstream_restrictions: BitstreamRestrictions::read(r)?, + bitstream_restrictions: BitstreamRestrictions::read(r, sps) + .map_err(SpsError::Field)?, }) } else { None @@ -877,10 +975,12 @@ pub struct SeqParameterSet { impl SeqParameterSet { pub fn from_bits(mut r: R) -> Result { let profile_idc = r.read_u8(8, "profile_idc")?.into(); - let sps = SeqParameterSet { + let constraint_flags = r.read_u8(8, "constraint_flags")?.into(); + let level_idc = r.read_u8(8, "level_idc")?; + let mut sps = SeqParameterSet { profile_idc, - constraint_flags: r.read_u8(8, "constraint_flags")?.into(), - level_idc: r.read_u8(8, "level_idc")?, + constraint_flags, + level_idc, seq_parameter_set_id: ParamSetId::from_u32(r.read_ue("seq_parameter_set_id")?) .map_err(SpsError::BadSeqParamSetId)?, chroma_info: ChromaInfo::read(&mut r, profile_idc)?, @@ -894,8 +994,13 @@ impl SeqParameterSet { frame_mbs_flags: FrameMbsFlags::read(&mut r)?, direct_8x8_inference_flag: r.read_bool("direct_8x8_inference_flag")?, frame_cropping: FrameCropping::read(&mut r)?, - vui_parameters: VuiParameters::read(&mut r)?, + // read the basic SPS data without reading VUI parameters yet, since checks of the + // bitstream restriction fields within the VUI parameters will need access to the + // initial SPS data + vui_parameters: None, }; + let vui_parameters = VuiParameters::read(&mut r, &sps)?; + sps.vui_parameters = vui_parameters; r.finish_rbsp()?; Ok(sps) } @@ -1019,6 +1124,285 @@ impl SeqParameterSet { } } +struct LevelLimit { + max_mbps: u32, + max_fs: u32, + max_dpb_mbs: u32, + max_br: u32, + max_cpb: u32, + max_vmv_r: u32, + min_cr: u8, + max_mvs_per2mb: Option, +} + +lazy_static! { + // "Table A-1 – Level limits" from the spec + static ref LEVEL_LIMITS: std::collections::HashMap = { + let mut m = std::collections::HashMap::new(); + m.insert( + Level::L1, + LevelLimit { + max_mbps: 1485, + max_fs: 99, + max_dpb_mbs: 396, + max_br: 64, + max_cpb: 175, + max_vmv_r: 64, + min_cr: 2, + max_mvs_per2mb: None, + }, + ); + m.insert( + Level::L1_b, + LevelLimit { + max_mbps: 1485, + max_fs: 99, + max_dpb_mbs: 396, + max_br: 128, + max_cpb: 350, + max_vmv_r: 64, + min_cr: 2, + max_mvs_per2mb: None, + }, + ); + m.insert( + Level::L1_1, + LevelLimit { + max_mbps: 3000, + max_fs: 396, + max_dpb_mbs: 900, + max_br: 192, + max_cpb: 500, + max_vmv_r: 128, + min_cr: 2, + max_mvs_per2mb: None, + }, + ); + m.insert( + Level::L1_2, + LevelLimit { + max_mbps: 6000, + max_fs: 396, + max_dpb_mbs: 2376, + max_br: 384, + max_cpb: 1000, + max_vmv_r: 128, + min_cr: 2, + max_mvs_per2mb: None, + }, + ); + m.insert( + Level::L1_3, + LevelLimit { + max_mbps: 11880, + max_fs: 396, + max_dpb_mbs: 2376, + max_br: 768, + max_cpb: 2000, + max_vmv_r: 128, + min_cr: 2, + max_mvs_per2mb: None, + }, + ); + m.insert( + Level::L2, + LevelLimit { + max_mbps: 11880, + max_fs: 396, + max_dpb_mbs: 2376, + max_br: 2000, + max_cpb: 2000, + max_vmv_r: 128, + min_cr: 2, + max_mvs_per2mb: None, + }, + ); + m.insert( + Level::L2_1, + LevelLimit { + max_mbps: 19800, + max_fs: 792, + max_dpb_mbs: 4752, + max_br: 4000, + max_cpb: 4000, + max_vmv_r: 256, + min_cr: 2, + max_mvs_per2mb: None, + }, + ); + m.insert( + Level::L2_2, + LevelLimit { + max_mbps: 20250, + max_fs: 1620, + max_dpb_mbs: 8100, + max_br: 4000, + max_cpb: 4000, + max_vmv_r: 256, + min_cr: 2, + max_mvs_per2mb: None, + }, + ); + m.insert( + Level::L3, + LevelLimit { + max_mbps: 40500, + max_fs: 1620, + max_dpb_mbs: 8100, + max_br: 10000, + max_cpb: 10000, + max_vmv_r: 256, + min_cr: 2, + max_mvs_per2mb: NonZeroU8::new(32), + }, + ); + m.insert( + Level::L3_1, + LevelLimit { + max_mbps: 108000, + max_fs: 3600, + max_dpb_mbs: 18000, + max_br: 14000, + max_cpb: 14000, + max_vmv_r: 512, + min_cr: 4, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m.insert( + Level::L3_2, + LevelLimit { + max_mbps: 216000, + max_fs: 5120, + max_dpb_mbs: 20480, + max_br: 20000, + max_cpb: 20000, + max_vmv_r: 512, + min_cr: 4, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m.insert( + Level::L4, + LevelLimit { + max_mbps: 245760, + max_fs: 8192, + max_dpb_mbs: 32768, + max_br: 20000, + max_cpb: 25000, + max_vmv_r: 512, + min_cr: 4, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m.insert( + Level::L4_1, + LevelLimit { + max_mbps: 245760, + max_fs: 8192, + max_dpb_mbs: 32768, + max_br: 50000, + max_cpb: 62500, + max_vmv_r: 512, + min_cr: 2, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m.insert( + Level::L4_2, + LevelLimit { + max_mbps: 522240, + max_fs: 8704, + max_dpb_mbs: 34816, + max_br: 50000, + max_cpb: 62500, + max_vmv_r: 512, + min_cr: 2, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m.insert( + Level::L5, + LevelLimit { + max_mbps: 589824, + max_fs: 22080, + max_dpb_mbs: 110400, + max_br: 135000, + max_cpb: 135000, + max_vmv_r: 512, + min_cr: 2, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m.insert( + Level::L5_1, + LevelLimit { + max_mbps: 983040, + max_fs: 36864, + max_dpb_mbs: 184320, + max_br: 240000, + max_cpb: 240000, + max_vmv_r: 512, + min_cr: 2, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m.insert( + Level::L5_2, + LevelLimit { + max_mbps: 2073600, + max_fs: 36864, + max_dpb_mbs: 184320, + max_br: 240000, + max_cpb: 240000, + max_vmv_r: 512, + min_cr: 2, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m.insert( + Level::L6, + LevelLimit { + max_mbps: 4177920, + max_fs: 139264, + max_dpb_mbs: 696320, + max_br: 240000, + max_cpb: 240000, + max_vmv_r: 8192, + min_cr: 2, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m.insert( + Level::L6_1, + LevelLimit { + max_mbps: 8355840, + max_fs: 139264, + max_dpb_mbs: 696320, + max_br: 480000, + max_cpb: 480000, + max_vmv_r: 8192, + min_cr: 2, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m.insert( + Level::L6_2, + LevelLimit { + max_mbps: 16711680, + max_fs: 139264, + max_dpb_mbs: 696320, + max_br: 800000, + max_cpb: 800000, + max_vmv_r: 8192, + min_cr: 2, + max_mvs_per2mb: NonZeroU8::new(16), + }, + ); + m + }; +} + #[cfg(test)] mod test { use crate::rbsp::{self, decode_nal, BitReader};