avutil/buffer: Avoid allocation of AVBuffer when using buffer pool

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

Reviewed-by: James Almer <jamrial@gmail.com>
Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt@outlook.com>

2 files changed, 41 insertions, 14 deletions

diff --git a/libavutil/buffer.c b/libavutil/buffer.c
index a32b7deb89..54590be566 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 {