[FFmpeg-devel] [PATCH] avcodec/cbs: Allocate more CodedBitstreamUnit at once in cbs_insert_unit()

[James Almer]

On 4/10/2020 9:00 PM, James Almer wrote:
> On 4/10/2020 8:53 PM, Michael Niedermayer wrote:
>> On Fri, Apr 10, 2020 at 05:44:25PM -0300, James Almer wrote:
>>> On 4/10/2020 5:23 PM, Michael Niedermayer wrote:
>>>> Fixes: Timeout (85sec -> 0.5sec)
>>>> 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: Michael Niedermayer <michael at niedermayer.cc>
>>>> ---
>>>>  libavcodec/cbs.c | 4 ++--
>>>>  1 file changed, 2 insertions(+), 2 deletions(-)
>>>>
>>>> diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
>>>> index 0bd5e1ac5d..42cb9711fa 100644
>>>> --- a/libavcodec/cbs.c
>>>> +++ b/libavcodec/cbs.c
>>>> @@ -693,11 +693,11 @@ static int cbs_insert_unit(CodedBitstreamContext *ctx,
>>>>              memmove(units + position + 1, units + position,
>>>>                      (frag->nb_units - position) * sizeof(*units));
>>>>      } else {
>>>> -        units = av_malloc_array(frag->nb_units + 1, sizeof(*units));
>>>> +        units = av_malloc_array(frag->nb_units*2 + 1, sizeof(*units));
>>>>          if (!units)
>>>>              return AVERROR(ENOMEM);
>>>>  
>>>> -        ++frag->nb_units_allocated;
>>>> +        frag->nb_units_allocated = 2*frag->nb_units_allocated + 1;
>>>
>>> Use ff_fast_malloc(), please. This is quite ugly and the *2 undocumented
>>> and not obvious.
>>
>> not sure i understood your suggestion correctly but 
>> ff_fast_malloc() deallocates its buffer and clears the size on error, 
>> which cbs_insert_unit() does not do.
>> so using ff_fast_malloc() feels a bit like fitting a square peg in a round
>> hole. But it works too
>> is that what you had in mind ? (your comment sounded like something simpler ...)
>> so maybe iam missing something
> 
> No, after thinking about it i realized it was not the best option and i
> sent a follow up reply about it, but i guess it was too late. If you
> have to change the implementation of ff_fast_malloc() then it's clearly
> not what should be used here. I didn't intend you to do this much work.
> 
> av_fast_realloc() could work instead as i mentioned in the follow up
> reply, i think. Or porting this to AVTreeNode instead of a flat array,
> but that's also quite a bit of work and will still allocate one node per
> call, so your fuzzer cases may still timeout. So if av_fast_realloc() is
> also not an option then maybe increase the buffer by a
> small-but-bigger-than-1 amount of units instead of duplicating its size
> each call, which can get quite big pretty fast.

Here's a version using av_fast_realloc(). FATE passes. Does it solve the
timeout?

diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
index 0bd5e1ac5d..d6cb94589f 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 >= frag->nb_units_allocated) {
+        size_t new_size;
+        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 (av_size_mult(frag->nb_units + 1, sizeof(*units), &new_size))
             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;
     }

-    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..4833a8a3db 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.