[FFmpeg-devel] [PATCH 20/36] avcodec/hevc_mp4toannexb_bsf: Stop reallocating output buffer

[Andreas Rheinhardt]

Instead parse the input packet twice: Once to get the size of the output
packet to be allocated (and to check the input packet for consistency)
and once to actually copy the data.

Not reallocating the output buffer also means that one can now use
a PutByteContext for writing; it improves readability.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt at gmail.com>
---
 libavcodec/hevc_mp4toannexb_bsf.c | 43 +++++++++++++++++++------------
 1 file changed, 27 insertions(+), 16 deletions(-)

diff --git a/libavcodec/hevc_mp4toannexb_bsf.c b/libavcodec/hevc_mp4toannexb_bsf.c
index 634bfe9c46..3776d5b279 100644
--- a/libavcodec/hevc_mp4toannexb_bsf.c
+++ b/libavcodec/hevc_mp4toannexb_bsf.c
@@ -119,10 +119,9 @@ static int hevc_mp4toannexb_init(AVBSFContext *ctx)
 static int hevc_mp4toannexb_filter(AVBSFContext *ctx, AVPacket *out)
 {
     HEVCBSFContext *s = ctx->priv_data;
+    PutByteContext pb;
     AVPacket *in;
-    GetByteContext gb;
 
-    int got_irap = 0;
     int i, ret = 0;
 
     ret = ff_bsf_get_packet(ctx, &in);
@@ -135,12 +134,17 @@ static int hevc_mp4toannexb_filter(AVBSFContext *ctx, AVPacket *out)
         return 0;
     }
 
+    for (int pass = 0; pass < 2; pass++) {
+        GetByteContext gb;
+        uint64_t out_size = 0;
+        int got_irap = 0;
+
     bytestream2_init(&gb, in->data, in->size);
 
     while (bytestream2_get_bytes_left(&gb)) {
         uint32_t nalu_size = 0;
         int      nalu_type;
-        int is_irap, add_extradata, extra_size, prev_size;
+            int is_irap, add_extradata, extra_size;
 
         if (bytestream2_get_bytes_left(&gb) < s->length_size) {
             ret = AVERROR_INVALIDDATA;
@@ -162,21 +166,28 @@ static int hevc_mp4toannexb_filter(AVBSFContext *ctx, AVPacket *out)
         extra_size    = add_extradata * ctx->par_out->extradata_size;
         got_irap     |= is_irap;
 
-        if (FFMIN(INT_MAX, SIZE_MAX) < 4ULL + nalu_size + extra_size) {
-            ret = AVERROR_INVALIDDATA;
-            goto fail;
-        }
-
-        prev_size = out->size;
+            if (!pass) {
+                out_size += extra_size + 4ULL + nalu_size;
+                bytestream2_skipu(&gb, nalu_size);
+            } else {
+                if (extra_size)
+                    bytestream2_put_bufferu(&pb, ctx->par_out->extradata, extra_size);
+                bytestream2_put_be32u(&pb, 1);
+                bytestream2_copy_bufferu(&pb, &gb, nalu_size);
+            }
+    }
 
-        ret = av_grow_packet(out, 4 + nalu_size + extra_size);
-        if (ret < 0)
-            goto fail;
+        if (!pass) {
+            if (out_size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE) {
+                ret = AVERROR(ERANGE);
+                goto fail;
+            }
+            ret = av_new_packet(out, out_size);
+            if (ret < 0)
+                goto fail;
 
-        if (extra_size)
-            memcpy(out->data + prev_size, ctx->par_out->extradata, extra_size);
-        AV_WB32(out->data + prev_size + extra_size, 1);
-        bytestream2_get_bufferu(&gb, out->data + prev_size + 4 + extra_size, nalu_size);
+            bytestream2_init_writer(&pb, out->data, out->size);
+        }
     }
 
     ret = av_packet_copy_props(out, in);
-- 
2.20.1