From 0fbd030b5d81745b17ef43f50cafcabc1e2adede Mon Sep 17 00:00:00 2001
From: jeffqjiangNew <142832361+jeffqjiangNew@users.noreply.github.com>
Date: Fri, 15 Dec 2023 15:54:43 -0500
Subject: [PATCH] Fixed the random crash associated with undecodable RASL
 pictures. (#141)

* * rocDecode/HEVC: Fixed the random crash associated with undecodable RASL pictures.
  - When a CRA picutre occurs, all the pictures in the DPB are emptied. However, the following RASL (Random Access Skipped Leading) picture can still reference these emptied pictures, making them undecodeble.
  - We initialized the reference picture lists with (0xFF), representing invalid picture index. Normally these invalid indices are replaced by valid values during RPS decoding. However, on undecodable RASL pictures, we can not find an existing reference picture in DPB, leaving the invalid index untouched and resulting invalid indexing later on.
  - Now we initialize the reference picture lists with a valid value 0. Also we add index range check on reference buffer at VA-API layer to avoid invalid memory access.

* * rocDecode/HEVC: Fixed a typo in sampel app message.
---
 samples/videoDecode/videodecode.cpp        |  2 +-
 src/parser/hevc_parser.cpp                 | 12 ++++++------
 src/rocdecode/vaapi/vaapi_videodecoder.cpp |  6 +++++-
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/samples/videoDecode/videodecode.cpp b/samples/videoDecode/videodecode.cpp
index b0c61f1cde..63ee431dc6 100644
--- a/samples/videoDecode/videodecode.cpp
+++ b/samples/videoDecode/videodecode.cpp
@@ -238,7 +238,7 @@ int main(int argc, char **argv) {
                 if (memcmp(digest, ref_md5, 16) == 0) {
                     std::cout << "MD5 digest matches the reference MD5 digest: ";
                 } else {
-                    std::cout << "MD5 digest does not matche the reference MD5 digest: ";
+                    std::cout << "MD5 digest does not match the reference MD5 digest: ";
                 }
                 ref_md5_file.seekg(0, std::ios_base::beg);
                 ref_md5_file.getline(ref_md5_string, 33);
diff --git a/src/parser/hevc_parser.cpp b/src/parser/hevc_parser.cpp
index b9c31ab530..a60ba63ea8 100644
--- a/src/parser/hevc_parser.cpp
+++ b/src/parser/hevc_parser.cpp
@@ -2073,13 +2073,13 @@ void HEVCVideoParser::DecodeRps() {
         /*
         * RPS derivation and picture marking
         */
-        // Init as "no reference picture"
+        // Init to a valid index value to take care of undecodable RASL pictures, whose reference pictures are emptied after a CRA.
         for (i = 0; i < HEVC_MAX_NUM_REF_PICS; i++) {
-            ref_pic_set_st_curr_before_[i] = 0xFF;
-            ref_pic_set_st_curr_after_[i] = 0xFF;
-            ref_pic_set_st_foll_[i] = 0xFF;
-            ref_pic_set_lt_curr_[i] = 0xFF;
-            ref_pic_set_lt_foll_[i] = 0xFF;
+            ref_pic_set_st_curr_before_[i] = 0;
+            ref_pic_set_st_curr_after_[i] = 0;
+            ref_pic_set_st_foll_[i] = 0;
+            ref_pic_set_lt_curr_[i] = 0;
+            ref_pic_set_lt_foll_[i] = 0;
         }
 
         // Mark all in DPB as unused. We will mark them back while we go through the ref lists. The rest will be actually unused.
diff --git a/src/rocdecode/vaapi/vaapi_videodecoder.cpp b/src/rocdecode/vaapi/vaapi_videodecoder.cpp
index 80ad39b9fb..2b1cf7f7cd 100644
--- a/src/rocdecode/vaapi/vaapi_videodecoder.cpp
+++ b/src/rocdecode/vaapi/vaapi_videodecoder.cpp
@@ -190,7 +190,7 @@ rocDecStatus VaapiVideoDecoder::SubmitDecode(RocdecPicParams *pPicParams) {
     VASurfaceID curr_surface_id;
 
     // Get the surface id for the current picture, assuming 1:1 mapping between DPB and VAAPI decoded surfaces.
-    if (pPicParams->CurrPicIdx >= va_surface_ids_.size()) {
+    if (pPicParams->CurrPicIdx >= va_surface_ids_.size() || pPicParams->CurrPicIdx < 0) {
         ERR("CurrPicIdx exceeded the VAAPI surface pool limit.");
         return ROCDEC_INVALID_PARAMETER;
     }
@@ -202,6 +202,10 @@ rocDecStatus VaapiVideoDecoder::SubmitDecode(RocdecPicParams *pPicParams) {
             pPicParams->pic_params.hevc.cur_pic.PicIdx = curr_surface_id;
             for (int i = 0; i < 15; i++) {
                 if (pPicParams->pic_params.hevc.ref_frames[i].PicIdx != 0xFF) {
+                    if (pPicParams->pic_params.hevc.ref_frames[i].PicIdx >= va_surface_ids_.size() || pPicParams->pic_params.hevc.ref_frames[i].PicIdx < 0) {
+                        ERR("Reference frame index exceeded the VAAPI surface pool limit.");
+                        return ROCDEC_INVALID_PARAMETER;
+                    }
                     pPicParams->pic_params.hevc.ref_frames[i].PicIdx = va_surface_ids_[pPicParams->pic_params.hevc.ref_frames[i].PicIdx];
                 }
             }