hevc: Fix a53 caption extraction

Just realized my previous patch doesn't work quite right.  I uploaded a better
sample file that actually has visible captions to /incoming/hevc_cc.ts.  I
tested with that file doing hevc->x264 and it works.

This is basically an exact copy of the existing h264 logic.

Signed-off-by: Will Kelleher <wkelleher@gogoair.com>
Signed-off-by: Michael Niedermayer <michael@niedermayer.cc>

2 files changed, 18 insertions, 10 deletions

diff --git a/libavcodec/hevc.c b/libavcodec/hevc.c
index 1fa5283da3..ece36f8eb0 100644
--- a/libavcodec/hevc.c
+++ b/libavcodec/hevc.c
@@ -2573,6 +2573,7 @@ static int set_side_data(HEVCContext *s)
         if (sd)
             memcpy(sd->data, s->a53_caption, s->a53_caption_size);
         av_freep(&s->a53_caption);
+        s->a53_caption_size = 0;
         s->avctx->properties |= FF_CODEC_PROPERTY_CLOSED_CAPTIONS;
     }
 
diff --git a/libavcodec/hevc_sei.c b/libavcodec/hevc_sei.c
index e853067cc4..46cd06b364 100644
--- a/libavcodec/hevc_sei.c
+++ b/libavcodec/hevc_sei.c
@@ -151,7 +151,6 @@ static int decode_registered_user_data_closed_caption(HEVCContext *s, int size)
     int flag;
     int user_data_type_code;
     int cc_count;
-    int i;
 
     GetBitContext *gb = &s->HEVClc->gb;
 
@@ -170,20 +169,28 @@ static int decode_registered_user_data_closed_caption(HEVCContext *s, int size)
             size -= 2;
 
             if (cc_count && size >= cc_count * 3) {
-                av_freep(&s->a53_caption);
-                s->a53_caption_size = cc_count * 3;
-
-                s->a53_caption = av_malloc(s->a53_caption_size);
-                if (!s->a53_caption)
-                    return(AVERROR(ENOMEM));
-
-                for (i = 0; i < s->a53_caption_size; i++) {
-                    s->a53_caption[i++] = get_bits(gb, 8);
+                const uint64_t new_size = (s->a53_caption_size + cc_count
+                                           * UINT64_C(3));
+                int i, ret;
+
+                if (new_size > INT_MAX)
+                    return AVERROR(EINVAL);
+
+                /* Allow merging of the cc data from two fields. */
+                ret = av_reallocp(&s->a53_caption, new_size);
+                if (ret < 0)
+                    return ret;
+
+                for (i = 0; i < cc_count; i++) {
+                    s->a53_caption[s->a53_caption_size++] = get_bits(gb, 8);
+                    s->a53_caption[s->a53_caption_size++] = get_bits(gb, 8);
+                    s->a53_caption[s->a53_caption_size++] = get_bits(gb, 8);
                 }
                 skip_bits(gb, 8); // marker_bits
             }
         }
     } else {
+        int i;
         for (i = 0; i < size - 1; i++)
             skip_bits(gb, 8);
     }