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

[James Almer]

On 9/14/2021 7:16 AM, Andreas Rheinhardt wrote:
> 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)

How about calling it something more generic BUFFER_FLAG_POOL or 
BUFFER_FLAG_FROM_POOL, to signal that the buffer was allocated by the 
buffer pool API? It can be reused, like in 
av_buffer_pool_buffer_get_opaque() to ensure there will be no crashes if 
the user passes it an AVBufferRef that was not created by the buffer 
pool API, and instead return NULL.

>   
>   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 {
>