[FFmpeg-devel] [PATCH] avcodec/cbs_mpeg2: use existing buffer reference if available when splitting fragments

[James Almer]

Signed-off-by: James Almer <jamrial at gmail.com>
---
 libavcodec/cbs_mpeg2.c | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/libavcodec/cbs_mpeg2.c b/libavcodec/cbs_mpeg2.c
index bfb64a0851..086d08ed64 100644
--- a/libavcodec/cbs_mpeg2.c
+++ b/libavcodec/cbs_mpeg2.c
@@ -146,14 +146,17 @@ static int cbs_mpeg2_split_fragment(CodedBitstreamContext *ctx,
             unit_size = (end - 4) - (start - 1);
         }
 
-        unit_data = av_malloc(unit_size + AV_INPUT_BUFFER_PADDING_SIZE);
-        if (!unit_data)
-            return AVERROR(ENOMEM);
-        memcpy(unit_data, start - 1, unit_size);
-        memset(unit_data + unit_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
+        if (header) {
+            unit_data = av_malloc(unit_size + AV_INPUT_BUFFER_PADDING_SIZE);
+            if (!unit_data)
+                return AVERROR(ENOMEM);
+            memcpy(unit_data, start - 1, unit_size);
+            memset(unit_data + unit_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
+        } else
+            unit_data = (uint8_t *)start - 1;
 
         err = ff_cbs_insert_unit_data(ctx, frag, i, unit_type,
-                                      unit_data, unit_size, NULL);
+                                      unit_data, unit_size, frag->data_ref);
         if (err < 0) {
             av_freep(&unit_data);
             return err;
-- 
2.17.0