[FFmpeg-devel] [PATCH] avutil/buffer: Avoid allocation of AVBuffer when using buffer pool

[Andreas Rheinhardt]

Do this by putting an AVBuffer structure into BufferPoolEntry and
reuse it for all subsequent uses of said BufferPoolEntry.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt at outlook.com>
---
 libavutil/buffer.c          | 44 +++++++++++++++++++++++++------------
 libavutil/buffer_internal.h | 11 ++++++++++
 2 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/libavutil/buffer.c b/libavutil/buffer.c
index b13eeadffb..4d9ccf74b7 100644
--- a/libavutil/buffer.c
+++ b/libavutil/buffer.c
@@ -26,16 +26,11 @@
 #include "mem.h"
 #include "thread.h"
 
-AVBufferRef *av_buffer_create(uint8_t *data, size_t size,
-                              void (*free)(void *opaque, uint8_t *data),
-                              void *opaque, int flags)
+static AVBufferRef *buffer_create(AVBuffer *buf, uint8_t *data, size_t size,
+                                  void (*free)(void *opaque, uint8_t *data),
+                                  void *opaque, int flags)
 {
     AVBufferRef *ref = NULL;
-    AVBuffer    *buf = NULL;
-
-    buf = av_mallocz(sizeof(*buf));
-    if (!buf)
-        return NULL;
 
     buf->data     = data;
     buf->size     = size;
@@ -47,10 +42,8 @@ AVBufferRef *av_buffer_create(uint8_t *data, size_t size,
     buf->flags = flags;
 
     ref = av_mallocz(sizeof(*ref));
-    if (!ref) {
-        av_freep(&buf);
+    if (!ref)
         return NULL;
-    }
 
     ref->buffer = buf;
     ref->data   = data;
@@ -59,6 +52,23 @@ AVBufferRef *av_buffer_create(uint8_t *data, size_t size,
     return ref;
 }
 
+AVBufferRef *av_buffer_create(uint8_t *data, size_t size,
+                              void (*free)(void *opaque, uint8_t *data),
+                              void *opaque, int flags)
+{
+    AVBufferRef *ret;
+    AVBuffer *buf = av_mallocz(sizeof(*buf));
+    if (!buf)
+        return NULL;
+
+    ret = buffer_create(buf, data, size, free, opaque, flags);
+    if (!ret) {
+        av_free(buf);
+        return NULL;
+    }
+    return ret;
+}
+
 void av_buffer_default_free(void *opaque, uint8_t *data)
 {
     av_free(data);
@@ -117,8 +127,12 @@ static void buffer_replace(AVBufferRef **dst, AVBufferRef **src)
         av_freep(dst);
 
     if (atomic_fetch_sub_explicit(&b->refcount, 1, memory_order_acq_rel) == 1) {
+        /* b->free below might already free the structure containing *b,
+         * so we have to read the flag now to avoid use-after-free. */
+        int free_avbuffer = !(b->flags_internal & BUFFER_FLAG_NO_FREE);
         b->free(b->opaque, b->data);
-        av_freep(&b);
+        if (free_avbuffer)
+            av_free(b);
     }
 }
 
@@ -378,11 +392,13 @@ AVBufferRef *av_buffer_pool_get(AVBufferPool *pool)
     ff_mutex_lock(&pool->mutex);
     buf = pool->pool;
     if (buf) {
-        ret = av_buffer_create(buf->data, pool->size, pool_release_buffer,
-                               buf, 0);
+        memset(&buf->buffer, 0, sizeof(buf->buffer));
+        ret = buffer_create(&buf->buffer, buf->data, pool->size,
+                            pool_release_buffer, buf, 0);
         if (ret) {
             pool->pool = buf->next;
             buf->next = NULL;
+            buf->buffer.flags_internal |= BUFFER_FLAG_NO_FREE;
         }
     } else {
         ret = pool_alloc_buffer(pool);
diff --git a/libavutil/buffer_internal.h b/libavutil/buffer_internal.h
index 839dc05f8f..bdff1b5b32 100644
--- a/libavutil/buffer_internal.h
+++ b/libavutil/buffer_internal.h
@@ -30,6 +30,11 @@
  * The buffer was av_realloc()ed, so it is reallocatable.
  */
 #define BUFFER_FLAG_REALLOCATABLE (1 << 0)
+/**
+ * The AVBuffer structure is part of a larger structure
+ * and should not be freed.
+ */
+#define BUFFER_FLAG_NO_FREE       (1 << 1)
 
 struct AVBuffer {
     uint8_t *data; /**< data described by this buffer */
@@ -73,6 +78,12 @@ typedef struct BufferPoolEntry {
 
     AVBufferPool *pool;
     struct BufferPoolEntry *next;
+
+    /*
+     * An AVBuffer structure to (re)use as AVBuffer for subsequent uses
+     * of this BufferPoolEntry.
+     */
+    AVBuffer buffer;
 } BufferPoolEntry;
 
 struct AVBufferPool {
-- 
2.30.2