[FFmpeg-devel] [PATCH] avcodec/cbs: use av_fast_realloc() in cbs_insert_unit()

[James Almer]

Fixes: Timeout
Fixes: 20791/clusterfuzz-testcase-minimized-ffmpeg_BSF_AV1_FRAME_SPLIT_fuzzer-5659537719951360
Fixes: 21214/clusterfuzz-testcase-minimized-ffmpeg_BSF_MPEG2_METADATA_fuzzer-5165560875974656
Fixes: 21247/clusterfuzz-testcase-minimized-ffmpeg_BSF_H264_METADATA_fuzzer-5715175257931776

Found-by: continuous fuzzing process https://github.com/google/oss-fuzz/tree/master/projects/ffmpeg
Signed-off-by: James Almer <jamrial at gmail.com>
---
This can be ported to av_fast_realloc_array() once that's committed.

 libavcodec/cbs.c | 37 ++++++++++++++++---------------------
 libavcodec/cbs.h |  7 +++++++
 2 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
index 0bd5e1ac5d..48ed7b5f32 100644
--- a/libavcodec/cbs.c
+++ b/libavcodec/cbs.c
@@ -161,6 +161,7 @@ void ff_cbs_fragment_free(CodedBitstreamContext *ctx,
 
     av_freep(&frag->units);
     frag->nb_units_allocated = 0;
+    frag->unit_buffer_size = 0;
 }
 
 static int cbs_read_fragment_content(CodedBitstreamContext *ctx,
@@ -684,35 +685,29 @@ static int cbs_insert_unit(CodedBitstreamContext *ctx,
                            CodedBitstreamFragment *frag,
                            int position)
 {
-    CodedBitstreamUnit *units;
+    CodedBitstreamUnit *units = frag->units;
 
-    if (frag->nb_units < frag->nb_units_allocated) {
-        units = frag->units;
+    if (frag->nb_units_allocated < frag->nb_units + 1) {
+        int new_size = frag->nb_units_allocated + 1;
+        void *tmp;
 
-        if (position < frag->nb_units)
-            memmove(units + position + 1, units + position,
-                    (frag->nb_units - position) * sizeof(*units));
-    } else {
-        units = av_malloc_array(frag->nb_units + 1, sizeof(*units));
-        if (!units)
+        if (new_size >= INT_MAX / sizeof(*units))
             return AVERROR(ENOMEM);
 
-        ++frag->nb_units_allocated;
-
-        if (position > 0)
-            memcpy(units, frag->units, position * sizeof(*units));
+        tmp = av_fast_realloc(units, &frag->unit_buffer_size,
+                              new_size * sizeof(*units));
+        if (!tmp)
+            return AVERROR(ENOMEM);
 
-        if (position < frag->nb_units)
-            memcpy(units + position + 1, frag->units + position,
-                   (frag->nb_units - position) * sizeof(*units));
+        frag->units = units = tmp;
+        frag->nb_units_allocated = new_size;
     }
 
-    memset(units + position, 0, sizeof(*units));
+    if (position < frag->nb_units)
+        memmove(units + position + 1, units + position,
+                (frag->nb_units - position) * sizeof(*units));
 
-    if (units != frag->units) {
-        av_free(frag->units);
-        frag->units = units;
-    }
+    memset(units + position, 0, sizeof(*units));
 
     ++frag->nb_units;
 
diff --git a/libavcodec/cbs.h b/libavcodec/cbs.h
index 9ca1fbd609..3209a82608 100644
--- a/libavcodec/cbs.h
+++ b/libavcodec/cbs.h
@@ -153,6 +153,13 @@ typedef struct CodedBitstreamFragment {
      */
      int             nb_units_allocated;
 
+    /**
+     * Size of allocated unit buffer.
+     *
+     * Must always be > nb_units_allocated; designed for internal use by cbs.
+     */
+     unsigned int    unit_buffer_size;
+
     /**
      * Pointer to an array of units of length nb_units_allocated.
      * Only the first nb_units are valid.
-- 
2.26.0