fix(rocDecode): Bound the shared NAL unit and SEI parsing in the base parser

projects/rocdecode/src/parser/avc_parser.cpp

@@ -125,13 +125,16 @@ ParserResult AvcVideoParser::ParsePictureData(const uint8_t *p_stream, uint32_t
         // Parse the NAL unit
         if (nal_unit_size_ >= 4) {
             // start code + NAL unit header = 4 bytes
-            int ebsp_size = nal_unit_size_ - 4 > RBSP_BUF_SIZE ? RBSP_BUF_SIZE : nal_unit_size_ - 4; // only copy enough bytes for header parsing
+            uint32_t ebsp_size = nal_unit_size_ - 4 > RBSP_BUF_SIZE ? RBSP_BUF_SIZE : nal_unit_size_ - 4; // only copy enough bytes for header parsing
 
             nal_unit_header_ = ParseNalUnitHeader(pic_data_buffer_ptr_[curr_start_code_offset_ + 3]);
             switch (nal_unit_header_.nal_unit_type) {
                 case kAvcNalTypeSeq_Parameter_Set: {
                     memcpy(rbsp_buf_, (pic_data_buffer_ptr_ + curr_start_code_offset_ + 4), ebsp_size);
-                    rbsp_size_ = EbspToRbsp(rbsp_buf_, 0, ebsp_size);
+                    if (EbspToRbsp(rbsp_buf_, 0, ebsp_size, &rbsp_size_) != PARSER_OK) {
+                        ErrorLog(g_rocdec_logger, "This NAL unit is skipped.");
+                        break;
+                    }
                     if ((ret2 = ParseSps(rbsp_buf_, rbsp_size_)) != PARSER_OK) {
                         ErrorLog(g_rocdec_logger, "Error occurred in SPS parsing. This SPS NAL unit is skipped.");
                     }
@@ -140,7 +143,10 @@ ParserResult AvcVideoParser::ParsePictureData(const uint8_t *p_stream, uint32_t
 
                 case kAvcNalTypePic_Parameter_Set: {
                     memcpy(rbsp_buf_, (pic_data_buffer_ptr_ + curr_start_code_offset_ + 4), ebsp_size);
-                    rbsp_size_ = EbspToRbsp(rbsp_buf_, 0, ebsp_size);
+                    if (EbspToRbsp(rbsp_buf_, 0, ebsp_size, &rbsp_size_) != PARSER_OK) {
+                        ErrorLog(g_rocdec_logger, "This NAL unit is skipped.");
+                        break;
+                    }
                     if ((ret2 = ParsePps(rbsp_buf_, rbsp_size_)) != PARSER_OK) {
                         ErrorLog(g_rocdec_logger, "Error occurred in PPS parsing. This PPS NAL unit is skipped.");
                     }
@@ -155,7 +161,10 @@ ParserResult AvcVideoParser::ParsePictureData(const uint8_t *p_stream, uint32_t
                     // Parse the slice header into a temporary first, so we can decide whether this
                     // slice begins a new primary coded picture before committing it.
                     memcpy(rbsp_buf_, (pic_data_buffer_ptr_ + curr_start_code_offset_ + 4), ebsp_size);
-                    rbsp_size_ = EbspToRbsp(rbsp_buf_, 0, ebsp_size);
+                    if (EbspToRbsp(rbsp_buf_, 0, ebsp_size, &rbsp_size_) != PARSER_OK) {
+                        ErrorLog(g_rocdec_logger, "This NAL unit is skipped.");
+                        break;
+                    }
                     AvcSliceHeader curr_slice_header;
                     if ((ret2 = ParseSliceHeader(rbsp_buf_, rbsp_size_, &curr_slice_header)) != PARSER_OK) {
                         ErrorLog(g_rocdec_logger, "Error occurred in slice header parsing. This slice NAL unit is skipped.");
@@ -252,7 +261,7 @@ ParserResult AvcVideoParser::ParsePictureData(const uint8_t *p_stream, uint32_t
 
                 case kAvcNalTypeSEI_Info: {
                     if (pfn_get_sei_message_cb_) {
-                        int sei_ebsp_size = nal_unit_size_ - 4; // copy the entire NAL unit
+                        uint32_t sei_ebsp_size = nal_unit_size_ - 4; // copy the entire NAL unit
                         if (sei_rbsp_buf_) {
                             if (sei_ebsp_size > sei_rbsp_buf_size_) {
                                 delete [] sei_rbsp_buf_;
@@ -264,7 +273,10 @@ ParserResult AvcVideoParser::ParsePictureData(const uint8_t *p_stream, uint32_t
                             sei_rbsp_buf_ = new uint8_t [sei_rbsp_buf_size_];
                         }
                         memcpy(sei_rbsp_buf_, (pic_data_buffer_ptr_ + curr_start_code_offset_ + 4), sei_ebsp_size);
-                        rbsp_size_ = EbspToRbsp(sei_rbsp_buf_, 0, sei_ebsp_size);
+                        if (EbspToRbsp(sei_rbsp_buf_, 0, sei_ebsp_size, &rbsp_size_) != PARSER_OK) {
+                            ErrorLog(g_rocdec_logger, "This NAL unit is skipped.");
+                            break;
+                        }
                         ParseSeiMessage(sei_rbsp_buf_, rbsp_size_);
                     }
                     break;

projects/rocdecode/src/parser/hevc_parser.cpp

@@ -589,27 +589,36 @@ ParserResult HevcVideoParser::ParsePictureData(const uint8_t* p_stream, uint32_t
         // Parse the NAL unit
         if (nal_unit_size_ >= 5) {
             // start code + NAL unit header = 5 bytes
-            int ebsp_size = nal_unit_size_ - 5 > RBSP_BUF_SIZE ? RBSP_BUF_SIZE : nal_unit_size_ - 5; // only copy enough bytes for header parsing
+            uint32_t ebsp_size = nal_unit_size_ - 5 > RBSP_BUF_SIZE ? RBSP_BUF_SIZE : nal_unit_size_ - 5; // only copy enough bytes for header parsing
 
             nal_unit_header_ = ParseNalUnitHeader(&pic_data_buffer_ptr_[curr_start_code_offset_ + 3]);
             switch (nal_unit_header_.nal_unit_type) {
                 case NAL_UNIT_VPS: {
                     memcpy(rbsp_buf_, (pic_data_buffer_ptr_ + curr_start_code_offset_ + 5), ebsp_size);
-                    rbsp_size_ = EbspToRbsp(rbsp_buf_, 0, ebsp_size);
+                    if (EbspToRbsp(rbsp_buf_, 0, ebsp_size, &rbsp_size_) != PARSER_OK) {
+                        ErrorLog(g_rocdec_logger, "This NAL unit is skipped.");
+                        break;
+                    }
                     ParseVps(rbsp_buf_, rbsp_size_);
                     break;
                 }
 
                 case NAL_UNIT_SPS: {
                     memcpy(rbsp_buf_, (pic_data_buffer_ptr_ + curr_start_code_offset_ + 5), ebsp_size);
-                    rbsp_size_ = EbspToRbsp(rbsp_buf_, 0, ebsp_size);
+                    if (EbspToRbsp(rbsp_buf_, 0, ebsp_size, &rbsp_size_) != PARSER_OK) {
+                        ErrorLog(g_rocdec_logger, "This NAL unit is skipped.");
+                        break;
+                    }
                     ParseSps(rbsp_buf_, rbsp_size_);
                     break;
                 }
 
                 case NAL_UNIT_PPS: {
                     memcpy(rbsp_buf_, (pic_data_buffer_ptr_ + curr_start_code_offset_ + 5), ebsp_size);
-                    rbsp_size_ = EbspToRbsp(rbsp_buf_, 0, ebsp_size);
+                    if (EbspToRbsp(rbsp_buf_, 0, ebsp_size, &rbsp_size_) != PARSER_OK) {
+                        ErrorLog(g_rocdec_logger, "This NAL unit is skipped.");
+                        break;
+                    }
                     ParsePps(rbsp_buf_, rbsp_size_);
                     break;
                 }
@@ -639,7 +648,10 @@ ParserResult HevcVideoParser::ParsePictureData(const uint8_t* p_stream, uint32_t
                     }
 
                     memcpy(rbsp_buf_, (pic_data_buffer_ptr_ + curr_start_code_offset_ + 5), ebsp_size);
-                    rbsp_size_ = EbspToRbsp(rbsp_buf_, 0, ebsp_size);
+                    if (EbspToRbsp(rbsp_buf_, 0, ebsp_size, &rbsp_size_) != PARSER_OK) {
+                        ErrorLog(g_rocdec_logger, "This NAL unit is skipped.");
+                        break;
+                    }
                     HevcSliceSegHeader *p_slice_header = &slice_info_list_[num_slices_].slice_header;
                     if ((ret2 = ParseSliceHeader(rbsp_buf_, rbsp_size_, p_slice_header)) != PARSER_OK) {
                         // we got an error while parsing this NAL unit. ignore and continue with next NAL unit
@@ -717,7 +729,7 @@ ParserResult HevcVideoParser::ParsePictureData(const uint8_t* p_stream, uint32_t
                 case NAL_UNIT_PREFIX_SEI:
                 case NAL_UNIT_SUFFIX_SEI: {
                     if (pfn_get_sei_message_cb_) {
-                        int sei_ebsp_size = nal_unit_size_ - 5; // copy the entire NAL unit
+                        uint32_t sei_ebsp_size = nal_unit_size_ - 5; // copy the entire NAL unit
                         if (sei_rbsp_buf_) {
                             if (sei_ebsp_size > sei_rbsp_buf_size_) {
                                 delete [] sei_rbsp_buf_;
@@ -729,7 +741,10 @@ ParserResult HevcVideoParser::ParsePictureData(const uint8_t* p_stream, uint32_t
                             sei_rbsp_buf_ = new uint8_t [sei_rbsp_buf_size_];
                         }
                         memcpy(sei_rbsp_buf_, (pic_data_buffer_ptr_ + curr_start_code_offset_ + 5), sei_ebsp_size);
-                        rbsp_size_ = EbspToRbsp(sei_rbsp_buf_, 0, sei_ebsp_size);
+                        if (EbspToRbsp(sei_rbsp_buf_, 0, sei_ebsp_size, &rbsp_size_) != PARSER_OK) {
+                            ErrorLog(g_rocdec_logger, "This NAL unit is skipped.");
+                            break;
+                        }
                         ParseSeiMessage(sei_rbsp_buf_, rbsp_size_);
                     }
                     break;

projects/rocdecode/src/parser/roc_video_parser.cpp

@@ -130,8 +130,10 @@ ParserResult RocVideoParser::GetNalUnit() {
     nal_unit_size_ = 0;
     curr_start_code_offset_ = next_start_code_offset_;  // save the current start code offset
 
-    // Search for the next start code
-    while (curr_byte_offset_ < pic_data_size_ - 2) {
+    // Search for the next start code. Both operands are uint32_t, so the size needs testing
+    // separately: a one byte packet would otherwise make pic_data_size_ - 2 wrap to 0xFFFFFFFF
+    // and run the scan past the end of the buffer.
+    while (pic_data_size_ >= 3 && curr_byte_offset_ < pic_data_size_ - 2) {
         if (pic_data_buffer_ptr_[curr_byte_offset_] == 0 && pic_data_buffer_ptr_[curr_byte_offset_ + 1] == 0 && pic_data_buffer_ptr_[curr_byte_offset_ + 2] == 0x01) {
             curr_start_code_offset_ = next_start_code_offset_;  // save the current start code offset
 
@@ -156,6 +158,15 @@ ParserResult RocVideoParser::GetNalUnit() {
         // No NAL unit in the frame data
         return PARSER_NOT_FOUND;
     }
+    // The subtractions below are unsigned, so an end offset below the start would yield a huge
+    // size rather than a negative one, and the nal_unit_size_ floor the callers apply before
+    // copying out of the NAL unit would not catch it. Check the ordering here instead.
+    if (curr_start_code_offset_ > pic_data_size_ ||
+        (start_code_found && next_start_code_offset_ < curr_start_code_offset_)) {
+        ErrorLog(g_rocdec_logger, "Start code offsets are out of order for the current picture.");
+        nal_unit_size_ = 0;
+        return PARSER_NOT_FOUND;
+    }
     if (start_code_found) {
         nal_unit_size_ = next_start_code_offset_ - curr_start_code_offset_;
         return PARSER_OK;
@@ -165,10 +176,12 @@ ParserResult RocVideoParser::GetNalUnit() {
     }
 }
 
-size_t RocVideoParser::EbspToRbsp(uint8_t *streamBuffer,size_t begin_bytepos, size_t end_bytepos) {
+ParserResult RocVideoParser::EbspToRbsp(uint8_t *streamBuffer,size_t begin_bytepos, size_t end_bytepos, size_t *p_rbsp_size) {
     int count = 0;
+    *p_rbsp_size = 0;
     if (end_bytepos < begin_bytepos) {
-        return end_bytepos;
+        *p_rbsp_size = end_bytepos;
+        return PARSER_OK;
     }
     uint8_t *streamBuffer_i = streamBuffer + begin_bytepos;
     uint8_t *streamBuffer_end = streamBuffer + end_bytepos;
@@ -181,10 +194,12 @@ size_t RocVideoParser::EbspToRbsp(uint8_t *streamBuffer,size_t begin_bytepos, si
             if (tmp == 0x03) {
                 //check the 4th uint8_t after 0x000003, except when cabac_zero_word is used, in which case the last three bytes of this NAL unit must be 0x000003
                 if ((streamBuffer_i + 1 != streamBuffer_end) && (streamBuffer_i[1] > 0x03)) {
-                    return static_cast<size_t>(-1);
+                    ErrorLog(g_rocdec_logger, "Malformed emulation prevention sequence in the NAL unit.");
+                    return PARSER_INVALID_ARG;
                 }
                 //if cabac_zero_word is used, the final uint8_t of this NAL unit(0x03) is discarded, and the last two bytes of RBSP must be 0x0000
                 if (streamBuffer_i + 1 == streamBuffer_end) {
+                    reduce_count++;  // discarded as well, so it is not part of the RBSP
                     break;
                 }
                 memmove(streamBuffer_i, streamBuffer_i + 1, streamBuffer_end-streamBuffer_i - 1);
@@ -202,31 +217,54 @@ size_t RocVideoParser::EbspToRbsp(uint8_t *streamBuffer,size_t begin_bytepos, si
         }
         streamBuffer_i++;
     }
-    return end_bytepos - begin_bytepos + reduce_count;
+    // Every discarded byte shortens the data, so the RBSP is the EBSP less reduce_count. Adding
+    // it instead reported more than the EBSP ever held: 00 00 03 01 is 3 bytes of RBSP but was
+    // reported as 5. For a full rbsp_buf_ that told the parse functions the buffer was larger
+    // than it is, which is exactly the kind of bound the rest of this change relies on.
+    *p_rbsp_size = end_bytepos - (begin_bytepos + reduce_count);
+    return PARSER_OK;
 }
 
-void RocVideoParser::ParseSeiMessage(uint8_t *nalu, size_t size) {
-    int offset = 0; // byte offset
-    int payload_type;
-    int payload_size;
+ParserResult RocVideoParser::ParseSeiMessage(uint8_t *nalu, size_t size) {
+    size_t offset = 0; // byte offset
+    // Accumulated in size_t so that a long run of ff_bytes cannot wrap the running total before
+    // it is range checked below.
+    size_t payload_type;
+    size_t payload_size;
 
+    // SEI is supplemental and does not affect the decode, so a message running past the end of
+    // the NAL unit stops the parse here and keeps whatever was read cleanly. The result is
+    // reported for the record; the callers log it and carry on with the picture.
     do {
         payload_type = 0;
-        while (nalu[offset] == 0xFF) {
+        while (offset < size && nalu[offset] == 0xFF) {
             payload_type += 255;  // ff_byte
             offset++;
         }
+        if (offset >= size) {
+            ErrorLog(g_rocdec_logger, "SEI payload type extends past the end of the NAL unit.");
+            return PARSER_OUT_OF_RANGE;
+        }
         payload_type += nalu[offset];  // last_payload_type_byte
         offset++;
 
         payload_size = 0;
-        while (nalu[offset] == 0xFF) {
+        while (offset < size && nalu[offset] == 0xFF) {
             payload_size += 255;  // ff_byte
             offset++;
         }
+        if (offset >= size) {
+            ErrorLog(g_rocdec_logger, "SEI payload size extends past the end of the NAL unit.");
+            return PARSER_OUT_OF_RANGE;
+        }
         payload_size += nalu[offset];  // last_payload_size_byte
         offset++;
 
+        if (payload_size > size - offset) {
+            ErrorLog(g_rocdec_logger, "SEI payload size (" + ROCDEC_TOSTR(payload_size) + ") exceeds the " + ROCDEC_TOSTR(size - offset) + " bytes left in the NAL unit.");
+            return PARSER_OUT_OF_RANGE;
+        }
+
         // We start with INIT_SEI_MESSAGE_COUNT. Should be enough for normal use cases. If not, resize.
         if((sei_message_count_ + 1) > sei_message_list_.size()) {
             sei_message_list_.resize((sei_message_count_ + 1));
@@ -236,7 +274,10 @@ void RocVideoParser::ParseSeiMessage(uint8_t *nalu, size_t size) {
 
         if (sei_payload_buf_) {
             if ((payload_size + sei_payload_size_) > sei_payload_buf_size_) {
-                uint8_t *tmp_ptr = new uint8_t [payload_size + sei_payload_size_];
+                // Record the new size: leaving sei_payload_buf_size_ at the old value made every
+                // later message compare against a stale bound and reallocate again.
+                sei_payload_buf_size_ = sei_payload_size_ + payload_size;
+                uint8_t *tmp_ptr = new uint8_t [sei_payload_buf_size_];
                 memcpy(tmp_ptr, sei_payload_buf_, sei_payload_size_); // save the existing payload
                 delete [] sei_payload_buf_;
                 sei_payload_buf_ = tmp_ptr;
@@ -254,4 +295,5 @@ void RocVideoParser::ParseSeiMessage(uint8_t *nalu, size_t size) {
 
         offset += payload_size;
     } while (offset < size && nalu[offset] != 0x80);
+    return PARSER_OK;
 }

projects/rocdecode/src/parser/roc_video_parser.h

@@ -184,11 +184,14 @@ class RocVideoParser {
 
     // NAL unit info
     int start_code_num_;              // number of start codes found so far
-    int curr_start_code_offset_;
-    int next_start_code_offset_;
-    int nal_unit_size_;
-
-    int                 rbsp_size_;
+    // Unsigned to match pic_data_size_ and curr_byte_offset_, so that none of the arithmetic
+    // between them converts signedness. GetNalUnit() range checks the two offsets before it
+    // derives nal_unit_size_ from them.
+    uint32_t curr_start_code_offset_;
+    uint32_t next_start_code_offset_;
+    uint32_t nal_unit_size_;
+
+    size_t              rbsp_size_; // size_t to match the parse functions it is passed to
     uint8_t             rbsp_buf_[RBSP_BUF_SIZE]; // to store parameter set or slice header RBSP
 
     int                 num_slices_;
@@ -220,20 +223,27 @@ class RocVideoParser {
     ParserResult GetNalUnit();
 
     /*! \brief Function to convert from Encapsulated Byte Sequence Packets to Raw Byte Sequence Payload
-     * 
+     *
      * \param [in,out] stream_buffer A pointer of <tt>uint8_t</tt> for the converted RBSP buffer.
      * \param [in] begin_bytepos Start position in the EBSP buffer to convert
      * \param [in] end_bytepos End position in the EBSP buffer to convert, generally it's size.
-     * \return Returns the size of the converted buffer in <tt>size_t</tt>
+     * \param [out] p_rbsp_size Size of the converted RBSP, which is the converted range less the
+     *             emulation prevention bytes that were discarded. Set to 0 on failure, so it can
+     *             never be mistaken for a length.
+     * \return <tt>ParserResult</tt>. PARSER_INVALID_ARG when an emulation prevention sequence is
+     *         malformed, in which case the caller should skip the NAL unit.
      */
-    size_t EbspToRbsp(uint8_t *stream_buffer, size_t begin_bytepos, size_t end_bytepos);
+    ParserResult EbspToRbsp(uint8_t *stream_buffer, size_t begin_bytepos, size_t end_bytepos, size_t *p_rbsp_size);
 
-    /*! \brief Function to parse Sei Message Info
+    /*! \brief Function to parse Sei Message Info. Stops at the first message that runs past the
+     *         end of the NAL unit, keeping the messages read before it. SEI does not affect the
+     *         decode, so callers are expected to note the result and carry on with the picture
+     *         rather than treat it as a picture level failure.
      * \param [in] nalu A pointer of <tt>uint8_t</tt> for the input stream to be parsed
      * \param [in] size Size of the input stream
-     * \return No return value
+     * \return <tt>ParserResult</tt>
      */
-    void ParseSeiMessage(uint8_t *nalu, size_t size);
+    ParserResult ParseSeiMessage(uint8_t *nalu, size_t size);
 
     /*! \brief Function to initialize the decoded buffer pool
      */