avcodec/hevc_ps: improve check for missing default display window bitstream

Fixes ticket #6644

Reviewed-by: Michael Niedermayer <michael@niedermayer.cc>
Signed-off-by: James Almer <jamrial@gmail.com>

1 files changed, 27 insertions, 6 deletions

diff --git a/libavcodec/hevc_ps.c b/libavcodec/hevc_ps.c
index ee31cc0..eb104ca 100644
--- a/libavcodec/hevc_ps.c
+++ b/libavcodec/hevc_ps.c
@@ -550,7 +550,7 @@ err:
 static void decode_vui(GetBitContext *gb, AVCodecContext *avctx,
                        int apply_defdispwin, HEVCSPS *sps)
 {
-    VUI *vui          = &sps->vui;
+    VUI backup_vui, *vui = &sps->vui;
     GetBitContext backup;
     int sar_present, alt = 0;
 
@@ -618,13 +618,14 @@ static void decode_vui(GetBitContext *gb, AVCodecContext *avctx,
     vui->field_seq_flag                = get_bits1(gb);
     vui->frame_field_info_present_flag = get_bits1(gb);
 
+    // Backup context in case an alternate header is detected
+    memcpy(&backup, gb, sizeof(backup));
+    memcpy(&backup_vui, vui, sizeof(backup_vui));
     if (get_bits_left(gb) >= 68 && show_bits_long(gb, 21) == 0x100000) {
         vui->default_display_window_flag = 0;
         av_log(avctx, AV_LOG_WARNING, "Invalid default display window\n");
     } else
         vui->default_display_window_flag = get_bits1(gb);
-    // Backup context in case an alternate header is detected
-    memcpy(&backup, gb, sizeof(backup));
 
     if (vui->default_display_window_flag) {
         int vert_mult  = 1 + (sps->chroma_format_idc < 2);
@@ -651,18 +652,19 @@ static void decode_vui(GetBitContext *gb, AVCodecContext *avctx,
         }
     }
 
+timing_info:
     vui->vui_timing_info_present_flag = get_bits1(gb);
 
     if (vui->vui_timing_info_present_flag) {
-        if( get_bits_left(gb) < 66) {
+        if( get_bits_left(gb) < 66 && !alt) {
             // The alternate syntax seem to have timing info located
             // at where def_disp_win is normally located
             av_log(avctx, AV_LOG_WARNING,
                    "Strange VUI timing information, retrying...\n");
-            vui->default_display_window_flag = 0;
-            memset(&vui->def_disp_win, 0, sizeof(vui->def_disp_win));
+            memcpy(vui, &backup_vui, sizeof(backup_vui));
             memcpy(gb, &backup, sizeof(backup));
             alt = 1;
+            goto timing_info;
         }
         vui->vui_num_units_in_tick               = get_bits_long(gb, 32);
         vui->vui_time_scale                      = get_bits_long(gb, 32);
@@ -680,6 +682,15 @@ static void decode_vui(GetBitContext *gb, AVCodecContext *avctx,
 
     vui->bitstream_restriction_flag = get_bits1(gb);
     if (vui->bitstream_restriction_flag) {
+        if (get_bits_left(gb) < 8 && !alt) {
+            av_log(avctx, AV_LOG_WARNING,
+                   "Strange VUI bitstream restriction information, retrying"
+                   " from timing information...\n");
+            memcpy(vui, &backup_vui, sizeof(backup_vui));
+            memcpy(gb, &backup, sizeof(backup));
+            alt = 1;
+            goto timing_info;
+        }
         vui->tiles_fixed_structure_flag              = get_bits1(gb);
         vui->motion_vectors_over_pic_boundaries_flag = get_bits1(gb);
         vui->restricted_ref_pic_lists_flag           = get_bits1(gb);
@@ -689,6 +700,16 @@ static void decode_vui(GetBitContext *gb, AVCodecContext *avctx,
         vui->log2_max_mv_length_horizontal           = get_ue_golomb_long(gb);
         vui->log2_max_mv_length_vertical             = get_ue_golomb_long(gb);
     }
+
+    if (get_bits_left(gb) < 1 && !alt) {
+        // XXX: Alternate syntax when sps_range_extension_flag != 0?
+        av_log(avctx, AV_LOG_WARNING,
+               "Overread in VUI, retrying from timing information...\n");
+        memcpy(vui, &backup_vui, sizeof(backup_vui));
+        memcpy(gb, &backup, sizeof(backup));
+        alt = 1;
+        goto timing_info;
+    }
 }
 
 static void set_default_scaling_list_data(ScalingList *sl)