[FFmpeg-devel] [PATCH v2] tools/target_dec_fuzzer: use refcounted packets

[James Almer]

Should reduce date copying considerably.

Signed-off-by: James Almer <jamrial at gmail.com>
---
Fixed a stupid mistake when checking the return value for av_new_packet().
Still untested.

 tools/target_dec_fuzzer.c | 71 ++++++++++++---------------------------
 1 file changed, 21 insertions(+), 50 deletions(-)

diff --git a/tools/target_dec_fuzzer.c b/tools/target_dec_fuzzer.c
index d83039417c..0d10503cfb 100644
--- a/tools/target_dec_fuzzer.c
+++ b/tools/target_dec_fuzzer.c
@@ -85,47 +85,6 @@ static int subtitle_handler(AVCodecContext *avctx, void *frame,
     return ret;
 }
 
-// Class to handle buffer allocation and resize for each frame
-typedef struct FuzzDataBuffer {
-    size_t size_;
-    uint8_t *data_;
-} FuzzDataBuffer;
-
-static void FDBCreate(FuzzDataBuffer *FDB) {
-    FDB->size_ = 0x1000;
-    FDB->data_ = av_malloc(FDB->size_);
-    if (!FDB->data_)
-        error("Failed memory allocation");
-}
-
-static void FDBDesroy(FuzzDataBuffer *FDB) { av_free(FDB->data_); }
-
-static void FDBRealloc(FuzzDataBuffer *FDB, size_t size) {
-    size_t needed = size + AV_INPUT_BUFFER_PADDING_SIZE;
-    av_assert0(needed > size);
-    if (needed > FDB->size_) {
-        av_free(FDB->data_);
-        FDB->size_ = needed;
-        FDB->data_ = av_malloc(FDB->size_);
-        if (!FDB->data_)
-            error("Failed memory allocation");
-    }
-}
-
-static void FDBPrepare(FuzzDataBuffer *FDB, AVPacket *dst, const uint8_t *data,
-                size_t size)
-{
-    FDBRealloc(FDB, size);
-    memcpy(FDB->data_, data, size);
-    size_t padd = FDB->size_ - size;
-    if (padd > AV_INPUT_BUFFER_PADDING_SIZE)
-        padd = AV_INPUT_BUFFER_PADDING_SIZE;
-    memset(FDB->data_ + size, 0, padd);
-    av_init_packet(dst);
-    dst->data = FDB->data_;
-    dst->size = size;
-}
-
 // Ensure we don't loop forever
 const uint32_t maxiteration = 8096;
 const uint64_t maxpixels_per_frame = 4096 * 4096;
@@ -135,7 +94,6 @@ static const uint64_t FUZZ_TAG = 0x4741542D5A5A5546ULL;
 
 int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
     const uint64_t fuzz_tag = FUZZ_TAG;
-    FuzzDataBuffer buffer;
     const uint8_t *last = data;
     const uint8_t *end = data + size;
     uint32_t it = 0;
@@ -186,6 +144,7 @@ int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
         error("Failed memory allocation");
 
     ctx->max_pixels = maxpixels_per_frame; //To reduce false positive OOM and hangs
+    ctx->refcounted_frames = 1;
 
     if (size > 1024) {
         GetByteContext gbc;
@@ -222,7 +181,6 @@ int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
     }
     parser_avctx->codec_id = ctx->codec_id;
 
-    FDBCreate(&buffer);
     int got_frame;
     AVFrame *frame = av_frame_alloc();
     if (!frame)
@@ -230,6 +188,7 @@ int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
 
     // Read very simple container
     AVPacket avpkt, parsepkt;
+    av_init_packet(&avpkt);
     while (data < end && it < maxiteration) {
         // Search for the TAG
         while (data + sizeof(fuzz_tag) < end) {
@@ -240,7 +199,10 @@ int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
         if (data + sizeof(fuzz_tag) > end)
             data = end;
 
-        FDBPrepare(&buffer, &parsepkt, last, data - last);
+        res = av_new_packet(&parsepkt, data - last);
+        if (res < 0)
+            error("Failed memory allocation");
+        memcpy(parsepkt.data, last, data - last);
         data += sizeof(fuzz_tag);
         last = data;
 
@@ -257,13 +219,21 @@ int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
                 avpkt.pts = parser->pts;
                 avpkt.dts = parser->dts;
                 avpkt.pos = parser->pos;
+                if (avpkt.data == parsepkt.data) {
+                    avpkt.buf = av_buffer_ref(parsepkt.buf);
+                    if (!avpkt.buf)
+                        error("Failed memory allocation");
+                } else {
+                    ret = av_packet_make_refcounted(&avpkt);
+                    if (ret < 0)
+                        error("Failed memory allocation");
+                }
                 if ( parser->key_frame == 1 ||
                     (parser->key_frame == -1 && parser->pict_type == AV_PICTURE_TYPE_I))
                     avpkt.flags |= AV_PKT_FLAG_KEY;
                 avpkt.flags |= parsepkt.flags & AV_PKT_FLAG_DISCARD;
             } else {
-                avpkt = parsepkt;
-                parsepkt.size = 0;
+                av_packet_move_ref(&avpkt, &parsepkt);
             }
 
           // Iterate through all data
@@ -284,16 +254,17 @@ int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
             avpkt.data += ret;
             avpkt.size -= ret;
           }
+          av_packet_unref(&avpkt);
         }
+        av_packet_unref(&parsepkt);
     }
 maximums_reached:
 
-    av_init_packet(&avpkt);
-    avpkt.data = NULL;
-    avpkt.size = 0;
+    av_packet_unref(&avpkt);
 
     do {
         got_frame = 0;
+        av_frame_unref(frame);
         decode_handler(ctx, frame, &got_frame, &avpkt);
     } while (got_frame == 1 && it++ < maxiteration);
 
@@ -303,6 +274,6 @@ maximums_reached:
     avcodec_free_context(&ctx);
     avcodec_free_context(&parser_avctx);
     av_parser_close(parser);
-    FDBDesroy(&buffer);
+    av_packet_unref(&parsepkt);
     return 0;
 }
-- 
2.22.0