From 3344bee39d45bcd809d190ce0cf217785446d538 Mon Sep 17 00:00:00 2001 From: Rajy Rawther Date: Mon, 11 Nov 2024 14:33:17 -0800 Subject: [PATCH] Added bitstream validation code for HEVC parser (#450) * fix for while loop hang * fix for while loop hang * add more data validation code in parser * minor change * addressed review comments * fix conformance failures with new code * added more checks [ROCm/rocdecode commit: 69ec7af8bdf26938cd7fa71df266fa2a14e1cb59] --- projects/rocdecode/src/parser/hevc_parser.cpp | 26 +++++++++++++++++-- projects/rocdecode/src/parser/hevc_parser.h | 1 + .../rocdecode/src/parser/roc_video_parser.cpp | 7 ++--- .../rocdecode/src/parser/roc_video_parser.h | 22 +++++++++++++--- 4 files changed, 47 insertions(+), 9 deletions(-) diff --git a/projects/rocdecode/src/parser/hevc_parser.cpp b/projects/rocdecode/src/parser/hevc_parser.cpp index 17bd9c46e7..a03f75eadc 100644 --- a/projects/rocdecode/src/parser/hevc_parser.cpp +++ b/projects/rocdecode/src/parser/hevc_parser.cpp @@ -565,7 +565,6 @@ ParserResult HevcVideoParser::ParsePictureData(const uint8_t* p_stream, uint32_t ERR(STR("Error: no start code found in the frame data.")); return ret; } - // Parse the NAL unit if (nal_unit_size_) { // start code + NAL unit header = 5 bytes @@ -625,7 +624,8 @@ ParserResult HevcVideoParser::ParsePictureData(const uint8_t* p_stream, uint32_t rbsp_size_ = EbspToRbsp(rbsp_buf_, 0, ebsp_size); HevcSliceSegHeader *p_slice_header = &slice_info_list_[num_slices_].slice_header; if ((ret2 = ParseSliceHeader(rbsp_buf_, rbsp_size_, p_slice_header)) != PARSER_OK) { - return ret2; + // we got an error while parsing this NAL unit. ignore and continue with next NAL unit + break; // ignore and continue to next nal_unit } // Start decode process @@ -1542,6 +1542,7 @@ ParserResult HevcVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size, HevcS // Set active VPS, SPS and PPS for the current slice m_active_pps_id_ = Parser::ExpGolomb::ReadUe(nalu, offset); + CHECK_ALLOWED_MAX(m_active_pps_id_, (MAX_PPS_COUNT - 1)); temp_sh.slice_pic_parameter_set_id = p_slice_header->slice_pic_parameter_set_id = m_active_pps_id_; pps_ptr = &m_pps_[m_active_pps_id_]; if ( pps_ptr->is_received == 0) { @@ -1603,6 +1604,7 @@ ParserResult HevcVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size, HevcS } int bits_slice_segment_address = (int)ceilf(log2f((float)pic_size_in_ctbs_y_)); temp_sh.slice_segment_address = p_slice_header->slice_segment_address = Parser::ReadBits(nalu, offset, bits_slice_segment_address); + //todo:: check for max slice_segment_address and error if exceeds max } if (!p_slice_header->dependent_slice_segment_flag) { @@ -1610,6 +1612,8 @@ ParserResult HevcVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size, HevcS p_slice_header->slice_reserved_flag[i] = Parser::GetBit(nalu, offset); } p_slice_header->slice_type = Parser::ExpGolomb::ReadUe(nalu, offset); + CHECK_ALLOWED_MAX(p_slice_header->slice_type, 2); + if (pps_ptr->output_flag_present_flag) { p_slice_header->pic_output_flag = Parser::GetBit(nalu, offset); } else { @@ -1650,6 +1654,10 @@ ParserResult HevcVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size, HevcS p_slice_header->num_long_term_sps = Parser::ExpGolomb::ReadUe(nalu, offset); } p_slice_header->num_long_term_pics = Parser::ExpGolomb::ReadUe(nalu, offset); + if (p_slice_header->num_long_term_pics > 16) { + ERR("Parsed Value exceeds allowed max"); + return PARSER_OUT_OF_RANGE; + } int bits_for_ltrp_in_sps = 0; while (sps_ptr->num_long_term_ref_pics_sps > (1 << bits_for_ltrp_in_sps)) { @@ -1702,8 +1710,10 @@ ParserResult HevcVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size, HevcS p_slice_header->num_ref_idx_active_override_flag = Parser::GetBit(nalu, offset); if (p_slice_header->num_ref_idx_active_override_flag) { p_slice_header->num_ref_idx_l0_active_minus1 = Parser::ExpGolomb::ReadUe(nalu, offset); + CHECK_ALLOWED_MAX(p_slice_header->num_ref_idx_l0_active_minus1, 14); if (p_slice_header->slice_type == HEVC_SLICE_TYPE_B) { p_slice_header->num_ref_idx_l1_active_minus1 = Parser::ExpGolomb::ReadUe(nalu, offset); + CHECK_ALLOWED_MAX(p_slice_header->num_ref_idx_l1_active_minus1, 14); } } else { p_slice_header->num_ref_idx_l0_active_minus1 = pps_ptr->num_ref_idx_l0_default_active_minus1; @@ -1746,6 +1756,7 @@ ParserResult HevcVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size, HevcS if (p_slice_header->ref_pic_list_modification_flag_l0) { for (int i = 0; i <= p_slice_header->num_ref_idx_l0_active_minus1; i++) { p_slice_header->list_entry_l0[i] = Parser::ReadBits(nalu, offset, list_entry_bits); + //todo;; check_max } } @@ -1754,6 +1765,7 @@ ParserResult HevcVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size, HevcS if (p_slice_header->ref_pic_list_modification_flag_l1) { for (int i = 0; i <= p_slice_header->num_ref_idx_l1_active_minus1; i++) { p_slice_header->list_entry_l1[i] = Parser::ReadBits(nalu, offset, list_entry_bits); + //todo::checkmax } } } @@ -1772,6 +1784,11 @@ ParserResult HevcVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size, HevcS } if ((p_slice_header->collocated_from_l0_flag && p_slice_header->num_ref_idx_l0_active_minus1 > 0) || (!p_slice_header->collocated_from_l0_flag && p_slice_header->num_ref_idx_l1_active_minus1 > 0)) { p_slice_header->collocated_ref_idx = Parser::ExpGolomb::ReadUe(nalu, offset); + if (p_slice_header->slice_type == HEVC_SLICE_TYPE_P || (p_slice_header->slice_type == HEVC_SLICE_TYPE_B && p_slice_header->collocated_from_l0_flag == 1)) { + CHECK_ALLOWED_MAX(p_slice_header->collocated_ref_idx, p_slice_header->num_ref_idx_l0_active_minus1); + } else { + CHECK_ALLOWED_MAX(p_slice_header->collocated_ref_idx, p_slice_header->num_ref_idx_l1_active_minus1); + } } } @@ -1779,12 +1796,15 @@ ParserResult HevcVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size, HevcS ParsePredWeightTable(p_slice_header, chroma_array_type, nalu, offset); } p_slice_header->five_minus_max_num_merge_cand = Parser::ExpGolomb::ReadUe(nalu, offset); + //CHECK_ALLOWED_MAX(p_slice_header->five_minus_max_num_merge_cand, 4); } p_slice_header->slice_qp_delta = Parser::ExpGolomb::ReadSe(nalu, offset); if (pps_ptr->pps_slice_chroma_qp_offsets_present_flag) { p_slice_header->slice_cb_qp_offset = Parser::ExpGolomb::ReadSe(nalu, offset); + CHECK_ALLOWED_RANGE(p_slice_header->slice_cb_qp_offset, -12, 12); p_slice_header->slice_cr_qp_offset = Parser::ExpGolomb::ReadSe(nalu, offset); + CHECK_ALLOWED_RANGE(p_slice_header->slice_cr_qp_offset, -12, 12); } if (pps_ptr->chroma_qp_offset_list_enabled_flag) { p_slice_header->cu_chroma_qp_offset_enabled_flag = Parser::GetBit(nalu, offset); @@ -1796,7 +1816,9 @@ ParserResult HevcVideoParser::ParseSliceHeader(uint8_t *nalu, size_t size, HevcS p_slice_header->slice_deblocking_filter_disabled_flag = Parser::GetBit(nalu, offset); if ( !p_slice_header->slice_deblocking_filter_disabled_flag ) { p_slice_header->slice_beta_offset_div2 = Parser::ExpGolomb::ReadSe(nalu, offset); + CHECK_ALLOWED_RANGE(p_slice_header->slice_beta_offset_div2, -6, 6); p_slice_header->slice_tc_offset_div2 = Parser::ExpGolomb::ReadSe(nalu, offset); + CHECK_ALLOWED_RANGE(p_slice_header->slice_tc_offset_div2, -6, 6); } } diff --git a/projects/rocdecode/src/parser/hevc_parser.h b/projects/rocdecode/src/parser/hevc_parser.h index f704733ef4..b0922e25d0 100644 --- a/projects/rocdecode/src/parser/hevc_parser.h +++ b/projects/rocdecode/src/parser/hevc_parser.h @@ -35,6 +35,7 @@ extern int scaling_list_default_1_2[2][6][64]; //size_id = 3 extern int scaling_list_default_3[1][2][64]; + class HevcVideoParser : public RocVideoParser { public: diff --git a/projects/rocdecode/src/parser/roc_video_parser.cpp b/projects/rocdecode/src/parser/roc_video_parser.cpp index 7db24a9870..e146a554b9 100644 --- a/projects/rocdecode/src/parser/roc_video_parser.cpp +++ b/projects/rocdecode/src/parser/roc_video_parser.cpp @@ -139,11 +139,12 @@ ParserResult RocVideoParser::GetNalUnit() { start_code_found = true; start_code_num_++; next_start_code_offset_ = curr_byte_offset_; - // Move the pointer 3 bytes forward - curr_byte_offset_ += 3; + // Move the pointer (3 + 2) bytes forward (start_code + nal header) + curr_byte_offset_ += 5; // For the very first NAL unit, search for the next start code (or reach the end of frame) - if (start_code_num_ == 1) { + // also detect empty nal units with no data + if (start_code_num_ == 1 ) { start_code_found = false; curr_start_code_offset_ = next_start_code_offset_; continue; diff --git a/projects/rocdecode/src/parser/roc_video_parser.h b/projects/rocdecode/src/parser/roc_video_parser.h index 6be47603f4..8555f18040 100644 --- a/projects/rocdecode/src/parser/roc_video_parser.h +++ b/projects/rocdecode/src/parser/roc_video_parser.h @@ -80,6 +80,21 @@ typedef struct { #define INIT_SEI_PAYLOAD_BUF_SIZE 1024 * 1024 // initial SEI payload buffer size, 1 MB #define DECODE_BUF_POOL_EXTENSION 2 +#define CHECK_ALLOWED_RANGE(val, min, max) { \ + if (val < min || val > max) { \ + ERR ("value not in range: " + TOSTR(val) + "allowed: " + TOSTR(min) + " " + TOSTR(max));\ + return PARSER_OUT_OF_RANGE; \ + } \ +} + +#define CHECK_ALLOWED_MAX(val, max) { \ + if (val > max) { \ + ERR ("value greater than maximum allowed value: " + TOSTR(val) + " max: " + TOSTR(max));\ + return PARSER_OUT_OF_RANGE; \ + } \ +} + + enum { kNotUsed = 0, kTopFieldUsedForDecode = 1, @@ -259,12 +274,11 @@ namespace Parser { return 0; // assert(0) } - uint32_t left_part = (0x1 << zero_bits_count) - 1; + uint32_t value = (0x1 << zero_bits_count) - 1; start_bit_idx++; - uint32_t rightPart = ReadBits(data, start_bit_idx, zero_bits_count); - return left_part + rightPart; + value += ReadBits(data, start_bit_idx, zero_bits_count); + return value; } - inline uint32_t ReadSe(const uint8_t *data, size_t &start_bit_idx) { uint32_t ue = ReadUe(data, start_bit_idx); // se From Ue