cbs: Refcount all the things!

This makes it easier for users of the CBS API to get alloc/free right -
all subelements use the buffer API so that it's clear how to free them.
It also allows eliding some redundant copies: the packet -> fragment copy
disappears after this change if the input packet is refcounted, and more
codec-specific cases are now possible (but not included in this patch).

15 files changed, 404 insertions, 279 deletions

diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
index a8d252f6c2..04ad2dfc41 100644
--- a/libavcodec/cbs.c
+++ b/libavcodec/cbs.c
@@ -21,6 +21,7 @@
 #include "config.h"
 
 #include "libavutil/avassert.h"
+#include "libavutil/buffer.h"
 #include "libavutil/common.h"
 
 #include "cbs.h"
@@ -95,11 +96,12 @@ void ff_cbs_close(CodedBitstreamContext **ctx_ptr)
 static void cbs_unit_uninit(CodedBitstreamContext *ctx,
                             CodedBitstreamUnit *unit)
 {
-    if (ctx->codec->free_unit && unit->content && !unit->content_external)
-        ctx->codec->free_unit(unit);
+    av_buffer_unref(&unit->content_ref);
+    unit->content = NULL;
 
-    av_freep(&unit->data);
-    unit->data_size = 0;
+    av_buffer_unref(&unit->data_ref);
+    unit->data             = NULL;
+    unit->data_size        = 0;
     unit->data_bit_padding = 0;
 }
 
@@ -113,7 +115,8 @@ void ff_cbs_fragment_uninit(CodedBitstreamContext *ctx,
     av_freep(&frag->units);
     frag->nb_units = 0;
 
-    av_freep(&frag->data);
+    av_buffer_unref(&frag->data_ref);
+    frag->data             = NULL;
     frag->data_size        = 0;
     frag->data_bit_padding = 0;
 }
@@ -133,6 +136,9 @@ static int cbs_read_fragment_content(CodedBitstreamContext *ctx,
                 continue;
         }
 
+        av_buffer_unref(&frag->units[i].content_ref);
+        frag->units[i].content = NULL;
+
         err = ctx->codec->read_unit(ctx, &frag->units[i]);
         if (err == AVERROR(ENOSYS)) {
             av_log(ctx->log_ctx, AV_LOG_VERBOSE,
@@ -169,6 +175,27 @@ int ff_cbs_read_extradata(CodedBitstreamContext *ctx,
     return cbs_read_fragment_content(ctx, frag);
 }
 
+static int cbs_fill_fragment_data(CodedBitstreamContext *ctx,
+                                  CodedBitstreamFragment *frag,
+                                  const uint8_t *data, size_t size)
+{
+    av_assert0(!frag->data && !frag->data_ref);
+
+    frag->data_ref =
+        av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
+    if (!frag->data_ref)
+        return AVERROR(ENOMEM);
+
+    frag->data      = frag->data_ref->data;
+    frag->data_size = size;
+
+    memcpy(frag->data, data, size);
+    memset(frag->data + size, 0,
+           AV_INPUT_BUFFER_PADDING_SIZE);
+
+    return 0;
+}
+
 int ff_cbs_read_packet(CodedBitstreamContext *ctx,
                        CodedBitstreamFragment *frag,
                        const AVPacket *pkt)
@@ -177,16 +204,24 @@ int ff_cbs_read_packet(CodedBitstreamContext *ctx,
 
     memset(frag, 0, sizeof(*frag));
 
-    frag->data      = pkt->data;
-    frag->data_size = pkt->size;
+    if (pkt->buf) {
+        frag->data_ref = av_buffer_ref(pkt->buf);
+        if (!frag->data_ref)
+            return AVERROR(ENOMEM);
+
+        frag->data      = pkt->data;
+        frag->data_size = pkt->size;
+
+    } else {
+        err = cbs_fill_fragment_data(ctx, frag, pkt->data, pkt->size);
+        if (err < 0)
+            return err;
+    }
 
     err = ctx->codec->split_fragment(ctx, frag, 0);
     if (err < 0)
         return err;
 
-    frag->data      = NULL;
-    frag->data_size = 0;
-
     return cbs_read_fragment_content(ctx, frag);
 }
 
@@ -198,17 +233,14 @@ int ff_cbs_read(CodedBitstreamContext *ctx,
 
     memset(frag, 0, sizeof(*frag));
 
-    // (We won't write to this during split.)
-    frag->data      = (uint8_t*)data;
-    frag->data_size = size;
+    err = cbs_fill_fragment_data(ctx, frag, data, size);
+    if (err < 0)
+        return err;
 
     err = ctx->codec->split_fragment(ctx, frag, 0);
     if (err < 0)
         return err;
 
-    frag->data      = NULL;
-    frag->data_size = 0;
-
     return cbs_read_fragment_content(ctx, frag);
 }
 
@@ -219,17 +251,25 @@ int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
     int err, i;
 
     for (i = 0; i < frag->nb_units; i++) {
-        if (!frag->units[i].content)
+        CodedBitstreamUnit *unit = &frag->units[i];
+
+        if (!unit->content)
             continue;
 
-        err = ctx->codec->write_unit(ctx, &frag->units[i]);
+        av_buffer_unref(&unit->data_ref);
+        unit->data = NULL;
+
+        err = ctx->codec->write_unit(ctx, unit);
         if (err < 0) {
             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to write unit %d "
-                   "(type %"PRIu32").\n", i, frag->units[i].type);
+                   "(type %"PRIu32").\n", i, unit->type);
             return err;
         }
     }
 
+    av_buffer_unref(&frag->data_ref);
+    frag->data = NULL;
+
     err = ctx->codec->assemble_fragment(ctx, frag);
     if (err < 0) {
         av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to assemble fragment.\n");
@@ -394,6 +434,45 @@ int ff_cbs_write_unsigned(CodedBitstreamContext *ctx, PutBitContext *pbc,
 }
 
 
+int ff_cbs_alloc_unit_content(CodedBitstreamContext *ctx,
+                              CodedBitstreamUnit *unit,
+                              size_t size,
+                              void (*free)(void *opaque, uint8_t *data))
+{
+    av_assert0(!unit->content && !unit->content_ref);
+
+    unit->content = av_mallocz(size);
+    if (!unit->content)
+        return AVERROR(ENOMEM);
+
+    unit->content_ref = av_buffer_create(unit->content, size,
+                                         free, ctx, 0);
+    if (!unit->content_ref) {
+        av_freep(&unit->content);
+        return AVERROR(ENOMEM);
+    }
+
+    return 0;
+}
+
+int ff_cbs_alloc_unit_data(CodedBitstreamContext *ctx,
+                           CodedBitstreamUnit *unit,
+                           size_t size)
+{
+    av_assert0(!unit->data && !unit->data_ref);
+
+    unit->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
+    if (!unit->data_ref)
+        return AVERROR(ENOMEM);
+
+    unit->data      = unit->data_ref->data;
+    unit->data_size = size;
+
+    memset(unit->data + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
+
+    return 0;
+}
+
 static int cbs_insert_unit(CodedBitstreamContext *ctx,
                            CodedBitstreamFragment *frag,
                            int position)
@@ -423,21 +502,35 @@ int ff_cbs_insert_unit_content(CodedBitstreamContext *ctx,
                                CodedBitstreamFragment *frag,
                                int position,
                                CodedBitstreamUnitType type,
-                               void *content)
+                               void *content,
+                               AVBufferRef *content_buf)
 {
+    CodedBitstreamUnit *unit;
+    AVBufferRef *content_ref;
     int err;
 
     if (position == -1)
         position = frag->nb_units;
     av_assert0(position >= 0 && position <= frag->nb_units);
 
+    if (content_buf) {
+        content_ref = av_buffer_ref(content_buf);
+        if (!content_ref)
+            return AVERROR(ENOMEM);
+    } else {
+        content_ref = NULL;
+    }
+
     err = cbs_insert_unit(ctx, frag, position);
-    if (err < 0)
+    if (err < 0) {
+        av_buffer_unref(&content_ref);
         return err;
+    }
 
-    frag->units[position].type             = type;
-    frag->units[position].content          = content;
-    frag->units[position].content_external = 1;
+    unit = &frag->units[position];
+    unit->type        = type;
+    unit->content     = content;
+    unit->content_ref = content_ref;
 
     return 0;
 }
@@ -446,21 +539,35 @@ int ff_cbs_insert_unit_data(CodedBitstreamContext *ctx,
                             CodedBitstreamFragment *frag,
                             int position,
                             CodedBitstreamUnitType type,
-                            uint8_t *data, size_t data_size)
+                            uint8_t *data, size_t data_size,
+                            AVBufferRef *data_buf)
 {
+    CodedBitstreamUnit *unit;
+    AVBufferRef *data_ref;
     int err;
 
     if (position == -1)
         position = frag->nb_units;
     av_assert0(position >= 0 && position <= frag->nb_units);
 
+    if (data_buf)
+        data_ref = av_buffer_ref(data_buf);
+    else
+        data_ref = av_buffer_create(data, data_size, NULL, NULL, 0);
+    if (!data_ref)
+        return AVERROR(ENOMEM);
+
     err = cbs_insert_unit(ctx, frag, position);
-    if (err < 0)
+    if (err < 0) {
+        av_buffer_unref(&data_ref);
         return err;
+    }
 
-    frag->units[position].type      = type;
-    frag->units[position].data      = data;
-    frag->units[position].data_size = data_size;
+    unit = &frag->units[position];
+    unit->type      = type;
+    unit->data      = data;
+    unit->data_size = data_size;
+    unit->data_ref  = data_ref;
 
     return 0;
 }
diff --git a/libavcodec/cbs.h b/libavcodec/cbs.h
index ffeca057ab..6505386708 100644
--- a/libavcodec/cbs.h
+++ b/libavcodec/cbs.h
@@ -22,6 +22,8 @@
 #include <stddef.h>
 #include <stdint.h>
 
+#include "libavutil/buffer.h"
+
 #include "avcodec.h"
 
 
@@ -81,6 +83,11 @@ typedef struct CodedBitstreamUnit {
      * This supports non-byte-aligned bitstreams.
      */
     size_t   data_bit_padding;
+    /**
+     * If data is reference counted, a reference to the buffer containing
+     * data.  Null if data is not reference counted.
+     */
+    AVBufferRef *data_ref;
 
     /**
      * Pointer to the decomposed form of this unit.
@@ -91,11 +98,10 @@ typedef struct CodedBitstreamUnit {
      */
     void *content;
     /**
-     * Whether the content was supplied externally.
-     *
-     * If so, it should not be freed when freeing the unit.
+     * If content is reference counted, a reference to the buffer containing
+     * content.  Null if content is not reference counted.
      */
-    int   content_external;
+    AVBufferRef *content_ref;
 } CodedBitstreamUnit;
 
 /**
@@ -123,6 +129,11 @@ typedef struct CodedBitstreamFragment {
      * The number of bits which should be ignored in the final byte.
      */
     size_t data_bit_padding;
+    /**
+     * If data is reference counted, a reference to the buffer containing
+     * data.  Null if data is not reference counted.
+     */
+    AVBufferRef *data_ref;
 
     /**
      * Number of units in this fragment.
@@ -279,27 +290,49 @@ void ff_cbs_fragment_uninit(CodedBitstreamContext *ctx,
 
 
 /**
+ * Allocate a new internal content buffer of the given size in the unit.
+ *
+ * The content will be zeroed.
+ */
+int ff_cbs_alloc_unit_content(CodedBitstreamContext *ctx,
+                              CodedBitstreamUnit *unit,
+                              size_t size,
+                              void (*free)(void *unit, uint8_t *content));
+
+/**
+ * Allocate a new internal data buffer of the given size in the unit.
+ *
+ * The data buffer will have input padding.
+ */
+int ff_cbs_alloc_unit_data(CodedBitstreamContext *ctx,
+                           CodedBitstreamUnit *unit,
+                           size_t size);
+
+/**
  * Insert a new unit into a fragment with the given content.
  *
- * The content structure continues to be owned by the caller, and
- * will not be freed when the unit is.
+ * The content structure continues to be owned by the caller if
+ * content_buf is not supplied.
  */
 int ff_cbs_insert_unit_content(CodedBitstreamContext *ctx,
                                CodedBitstreamFragment *frag,
                                int position,
                                CodedBitstreamUnitType type,
-                               void *content);
+                               void *content,
+                               AVBufferRef *content_buf);
 
 /**
  * Insert a new unit into a fragment with the given data bitstream.
  *
- * The data buffer will be owned by the unit after this operation.
+ * If data_buf is not supplied then data must have been allocated with
+ * av_malloc() and will become owned by the unit after this call.
  */
 int ff_cbs_insert_unit_data(CodedBitstreamContext *ctx,
                             CodedBitstreamFragment *frag,
                             int position,
                             CodedBitstreamUnitType type,
-                            uint8_t *data, size_t data_size);
+                            uint8_t *data, size_t data_size,
+                            AVBufferRef *data_buf);
 
 /**
  * Delete a unit from a fragment and free all memory it uses.
diff --git a/libavcodec/cbs_h264.h b/libavcodec/cbs_h264.h
index b58f19f17e..14ea69ae28 100644
--- a/libavcodec/cbs_h264.h
+++ b/libavcodec/cbs_h264.h
@@ -266,12 +266,14 @@ typedef struct H264RawSEIUserDataRegistered {
     uint8_t itu_t_t35_country_code_extension_byte;
     uint8_t *data;
     size_t data_length;
+    AVBufferRef *data_ref;
 } H264RawSEIUserDataRegistered;
 
 typedef struct H264RawSEIUserDataUnregistered {
     uint8_t uuid_iso_iec_11578[16];
     uint8_t *data;
     size_t data_length;
+    AVBufferRef *data_ref;
 } H264RawSEIUserDataUnregistered;
 
 typedef struct H264RawSEIRecoveryPoint {
@@ -304,6 +306,7 @@ typedef struct H264RawSEIPayload {
         struct {
             uint8_t *data;
             size_t data_length;
+            AVBufferRef *data_ref;
         } other;
     } payload;
 } H264RawSEIPayload;
@@ -399,6 +402,7 @@ typedef struct H264RawSlice {
     uint8_t *data;
     size_t   data_size;
     int      data_bit_start;
+    AVBufferRef *data_ref;
 } H264RawSlice;
 
 
diff --git a/libavcodec/cbs_h2645.c b/libavcodec/cbs_h2645.c
index 9d05d59156..5b23d6104b 100644
--- a/libavcodec/cbs_h2645.c
+++ b/libavcodec/cbs_h2645.c
@@ -291,9 +291,10 @@ static int cbs_h2645_read_more_rbsp_data(BitstreamContext *bc)
 #define byte_alignment(rw) (bitstream_tell(rw) % 8)
 
 #define allocate(name, size) do { \
-        name = av_mallocz(size); \
-        if (!name) \
+        name ## _ref = av_buffer_allocz(size); \
+        if (!name ## _ref) \
             return AVERROR(ENOMEM); \
+        name = name ## _ref->data; \
     } while (0)
 
 #define FUNC(name) FUNC_H264(READWRITE, name)
@@ -393,82 +394,68 @@ static int cbs_h2645_read_more_rbsp_data(BitstreamContext *bc)
 #undef allocate
 
 
-static void cbs_h264_free_sei(H264RawSEI *sei)
+static void cbs_h264_free_sei_payload(H264RawSEIPayload *payload)
 {
-    int i;
-    for (i = 0; i < sei->payload_count; i++) {
-        H264RawSEIPayload *payload = &sei->payload[i];
-
-        switch (payload->payload_type) {
-        case H264_SEI_TYPE_BUFFERING_PERIOD:
-        case H264_SEI_TYPE_PIC_TIMING:
-        case H264_SEI_TYPE_RECOVERY_POINT:
-        case H264_SEI_TYPE_DISPLAY_ORIENTATION:
-            break;
-        case H264_SEI_TYPE_USER_DATA_REGISTERED:
-            av_freep(&payload->payload.user_data_registered.data);
-            break;
-        case H264_SEI_TYPE_USER_DATA_UNREGISTERED:
-            av_freep(&payload->payload.user_data_unregistered.data);
-            break;
-        default:
-            av_freep(&payload->payload.other.data);
-            break;
-        }
+    switch (payload->payload_type) {
+    case H264_SEI_TYPE_BUFFERING_PERIOD:
+    case H264_SEI_TYPE_PIC_TIMING:
+    case H264_SEI_TYPE_RECOVERY_POINT:
+    case H264_SEI_TYPE_DISPLAY_ORIENTATION:
+        break;
+    case H264_SEI_TYPE_USER_DATA_REGISTERED:
+        av_buffer_unref(&payload->payload.user_data_registered.data_ref);
+        break;
+    case H264_SEI_TYPE_USER_DATA_UNREGISTERED:
+        av_buffer_unref(&payload->payload.user_data_unregistered.data_ref);
+        break;
+    default:
+        av_buffer_unref(&payload->payload.other.data_ref);
+        break;
     }
 }
 
-static void cbs_h264_free_slice(H264RawSlice *slice)
+static void cbs_h264_free_sei(void *unit, uint8_t *content)
 {
-    av_freep(&slice->data);
+    H264RawSEI *sei = (H264RawSEI*)content;
+    int i;
+    for (i = 0; i < sei->payload_count; i++)
+        cbs_h264_free_sei_payload(&sei->payload[i]);
+    av_freep(&content);
 }
 
-static void cbs_h264_free_nal_unit(CodedBitstreamUnit *unit)
+static void cbs_h264_free_slice(void *unit, uint8_t *content)
 {
-    switch (unit->type) {
-    case H264_NAL_SEI:
-        cbs_h264_free_sei(unit->content);
-        break;
-    case H264_NAL_IDR_SLICE:
-    case H264_NAL_SLICE:
-        cbs_h264_free_slice(unit->content);
-        break;
-    }
-    av_freep(&unit->content);
+    H264RawSlice *slice = (H264RawSlice*)content;
+    av_buffer_unref(&slice->data_ref);
+    av_freep(&content);
 }
 
-static void cbs_h265_free_nal_unit(CodedBitstreamUnit *unit)
+static void cbs_h265_free_vps(void *unit, uint8_t *content)
 {
-    switch (unit->type) {
-    case HEVC_NAL_VPS:
-        av_freep(&((H265RawVPS*)unit->content)->extension_data.data);
-        break;
-    case HEVC_NAL_SPS:
-        av_freep(&((H265RawSPS*)unit->content)->extension_data.data);
-        break;
-    case HEVC_NAL_PPS:
-        av_freep(&((H265RawPPS*)unit->content)->extension_data.data);
-        break;
-    case HEVC_NAL_TRAIL_N:
-    case HEVC_NAL_TRAIL_R:
-    case HEVC_NAL_TSA_N:
-    case HEVC_NAL_TSA_R:
-    case HEVC_NAL_STSA_N:
-    case HEVC_NAL_STSA_R:
-    case HEVC_NAL_RADL_N:
-    case HEVC_NAL_RADL_R:
-    case HEVC_NAL_RASL_N:
-    case HEVC_NAL_RASL_R:
-    case HEVC_NAL_BLA_W_LP:
-    case HEVC_NAL_BLA_W_RADL:
-    case HEVC_NAL_BLA_N_LP:
-    case HEVC_NAL_IDR_W_RADL:
-    case HEVC_NAL_IDR_N_LP:
-    case HEVC_NAL_CRA_NUT:
-        av_freep(&((H265RawSlice*)unit->content)->data);
-        break;
-    }
-    av_freep(&unit->content);
+    H265RawVPS *vps = (H265RawVPS*)content;
+    av_buffer_unref(&vps->extension_data.data_ref);
+    av_freep(&content);
+}
+
+static void cbs_h265_free_sps(void *unit, uint8_t *content)
+{
+    H265RawSPS *sps = (H265RawSPS*)content;
+    av_buffer_unref(&sps->extension_data.data_ref);
+    av_freep(&content);
+}
+
+static void cbs_h265_free_pps(void *unit, uint8_t *content)
+{
+    H265RawPPS *pps = (H265RawPPS*)content;
+    av_buffer_unref(&pps->extension_data.data_ref);
+    av_freep(&content);
+}
+
+static void cbs_h265_free_slice(void *unit, uint8_t *content)
+{
+    H265RawSlice *slice = (H265RawSlice*)content;
+    av_buffer_unref(&slice->data_ref);
+    av_freep(&content);
 }
 
 static int cbs_h2645_fragment_add_nals(CodedBitstreamContext *ctx,
@@ -494,7 +481,7 @@ static int cbs_h2645_fragment_add_nals(CodedBitstreamContext *ctx,
         memset(data + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 
         err = ff_cbs_insert_unit_data(ctx, frag, -1, nal->type,
-                                      data, nal->size);
+                                      data, nal->size, NULL);
         if (err < 0) {
             av_freep(&data);
             return err;
@@ -705,35 +692,32 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
         {
             H264RawSPS *sps;
 
-            sps = av_mallocz(sizeof(*sps));
-            if (!sps)
-                return AVERROR(ENOMEM);
+            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*sps), NULL);
+            if (err < 0)
+                return err;
+            sps = unit->content;
+
             err = cbs_h264_read_sps(ctx, &bc, sps);
-            if (err >= 0)
-                err = cbs_h264_replace_sps(ctx, sps);
-            if (err < 0) {
-                av_free(sps);
+            if (err < 0)
                 return err;
-            }
 
-            unit->content = sps;
+            err = cbs_h264_replace_sps(ctx, sps);
+            if (err < 0)
+                return err;
         }
         break;
 
     case H264_NAL_SPS_EXT:
         {
-            H264RawSPSExtension *sps_ext;
-
-            sps_ext = av_mallocz(sizeof(*sps_ext));
-            if (!sps_ext)
-                return AVERROR(ENOMEM);
-            err = cbs_h264_read_sps_extension(ctx, &bc, sps_ext);
-            if (err < 0) {
-                av_free(sps_ext);
+            err = ff_cbs_alloc_unit_content(ctx, unit,
+                                            sizeof(H264RawSPSExtension),
+                                            NULL);
+            if (err < 0)
                 return err;
-            }
 
-            unit->content = sps_ext;
+            err = cbs_h264_read_sps_extension(ctx, &bc, unit->content);
+            if (err < 0)
+                return err;
         }
         break;
 
@@ -741,18 +725,18 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
         {
             H264RawPPS *pps;
 
-            pps = av_mallocz(sizeof(*pps));
-            if (!pps)
-                return AVERROR(ENOMEM);
+            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*pps), NULL);
+            if (err < 0)
+                return err;
+            pps = unit->content;
+
             err = cbs_h264_read_pps(ctx, &bc, pps);
-            if (err >= 0)
-                err = cbs_h264_replace_pps(ctx, pps);
-            if (err < 0) {
-                av_free(pps);
+            if (err < 0)
                 return err;
-            }
 
-            unit->content = pps;
+            err = cbs_h264_replace_pps(ctx, pps);
+            if (err < 0)
+                return err;
         }
         break;
 
@@ -763,14 +747,15 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
             H264RawSlice *slice;
             int pos, len;
 
-            slice = av_mallocz(sizeof(*slice));
-            if (!slice)
-                return AVERROR(ENOMEM);
+            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*slice),
+                                            &cbs_h264_free_slice);
+            if (err < 0)
+                return err;
+            slice = unit->content;
+
             err = cbs_h264_read_slice_header(ctx, &bc, &slice->header);
-            if (err < 0) {
-                av_free(slice);
+            if (err < 0)
                 return err;
-            }
 
             pos = bitstream_tell(&bc);
             len = unit->data_size;
@@ -783,54 +768,42 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
             }
 
             slice->data_size = len - pos / 8;
-            slice->data = av_malloc(slice->data_size +
-                                    AV_INPUT_BUFFER_PADDING_SIZE);
-            if (!slice->data) {
-                av_free(slice);
+            slice->data_ref  = av_buffer_alloc(slice->data_size +
+                                               AV_INPUT_BUFFER_PADDING_SIZE);
+            if (!slice->data_ref)
                 return AVERROR(ENOMEM);
-            }
+            slice->data = slice->data_ref->data;
             memcpy(slice->data,
                    unit->data + pos / 8, slice->data_size);
             memset(slice->data + slice->data_size, 0,
                    AV_INPUT_BUFFER_PADDING_SIZE);
             slice->data_bit_start = pos % 8;
-
-            unit->content = slice;
         }
         break;
 
     case H264_NAL_AUD:
         {
-            H264RawAUD *aud;
-
-            aud = av_mallocz(sizeof(*aud));
-            if (!aud)
-                return AVERROR(ENOMEM);
-            err = cbs_h264_read_aud(ctx, &bc, aud);
-            if (err < 0) {
-                av_free(aud);
+            err = ff_cbs_alloc_unit_content(ctx, unit,
+                                            sizeof(H264RawAUD), NULL);
+            if (err < 0)
                 return err;
-            }
 
-            unit->content = aud;
+            err = cbs_h264_read_aud(ctx, &bc, unit->content);
+            if (err < 0)
+                return err;
         }
         break;
 
     case H264_NAL_SEI:
         {
-            H264RawSEI *sei;
-
-            sei = av_mallocz(sizeof(*sei));
-            if (!sei)
-                return AVERROR(ENOMEM);
-            err = cbs_h264_read_sei(ctx, &bc, sei);
-            if (err < 0) {
-                cbs_h264_free_sei(sei);
-                av_free(sei);
+            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(H264RawSEI),
+                                            &cbs_h264_free_sei);
+            if (err < 0)
                 return err;
-            }
 
-            unit->content = sei;
+            err = cbs_h264_read_sei(ctx, &bc, unit->content);
+            if (err < 0)
+                return err;
         }
         break;
 
@@ -856,36 +829,38 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
         {
             H265RawVPS *vps;
 
-            vps = av_mallocz(sizeof(*vps));
-            if (!vps)
-                return AVERROR(ENOMEM);
+            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*vps),
+                                            &cbs_h265_free_vps);
+            if (err < 0)
+                return err;
+            vps = unit->content;
+
             err = cbs_h265_read_vps(ctx, &bc, vps);
-            if (err >= 0)
-                err = cbs_h265_replace_vps(ctx, vps);
-            if (err < 0) {
-                av_free(vps);
+            if (err < 0)
                 return err;
-            }
 
-            unit->content = vps;
+            err = cbs_h265_replace_vps(ctx, vps);
+            if (err < 0)
+                return err;
         }
         break;
     case HEVC_NAL_SPS:
         {
             H265RawSPS *sps;
 
-            sps = av_mallocz(sizeof(*sps));
-            if (!sps)
-                return AVERROR(ENOMEM);
+            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*sps),
+                                            &cbs_h265_free_sps);
+            if (err < 0)
+                return err;
+            sps = unit->content;
+
             err = cbs_h265_read_sps(ctx, &bc, sps);
-            if (err >= 0)
-                err = cbs_h265_replace_sps(ctx, sps);
-            if (err < 0) {
-                av_free(sps);
+            if (err < 0)
                 return err;
-            }
 
-            unit->content = sps;
+            err = cbs_h265_replace_sps(ctx, sps);
+            if (err < 0)
+                return err;
         }
         break;
 
@@ -893,18 +868,19 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
         {
             H265RawPPS *pps;
 
-            pps = av_mallocz(sizeof(*pps));
-            if (!pps)
-                return AVERROR(ENOMEM);
+            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*pps),
+                                            &cbs_h265_free_pps);
+            if (err < 0)
+                return err;
+            pps = unit->content;
+
             err = cbs_h265_read_pps(ctx, &bc, pps);
-            if (err >= 0)
-                err = cbs_h265_replace_pps(ctx, pps);
-            if (err < 0) {
-                av_free(pps);
+            if (err < 0)
                 return err;
-            }
 
-            unit->content = pps;
+            err = cbs_h265_replace_pps(ctx, pps);
+            if (err < 0)
+                return err;
         }
         break;
 
@@ -928,14 +904,15 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
             H265RawSlice *slice;
             int pos, len;
 
-            slice = av_mallocz(sizeof(*slice));
-            if (!slice)
-                return AVERROR(ENOMEM);
+            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*slice),
+                                            &cbs_h265_free_slice);
+            if (err < 0)
+                return err;
+            slice = unit->content;
+
             err = cbs_h265_read_slice_segment_header(ctx, &bc, &slice->header);
-            if (err < 0) {
-                av_free(slice);
+            if (err < 0)
                 return err;
-            }
 
             pos = bitstream_tell(&bc);
             len = unit->data_size;
@@ -948,36 +925,29 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
             }
 
             slice->data_size = len - pos / 8;
-            slice->data = av_malloc(slice->data_size +
-                                    AV_INPUT_BUFFER_PADDING_SIZE);
-            if (!slice->data) {
-                av_free(slice);
+            slice->data_ref  = av_buffer_alloc(slice->data_size +
+                                               AV_INPUT_BUFFER_PADDING_SIZE);
+            if (!slice->data_ref)
                 return AVERROR(ENOMEM);
-            }
+            slice->data = slice->data_ref->data;
             memcpy(slice->data,
                    unit->data + pos / 8, slice->data_size);
             memset(slice->data + slice->data_size, 0,
                    AV_INPUT_BUFFER_PADDING_SIZE);
             slice->data_bit_start = pos % 8;
-
-            unit->content = slice;
         }
         break;
 
     case HEVC_NAL_AUD:
         {
-            H265RawAUD *aud;
-
-            aud = av_mallocz(sizeof(*aud));
-            if (!aud)
-                return AVERROR(ENOMEM);
-            err = cbs_h265_read_aud(ctx, &bc, aud);
-            if (err < 0) {
-                av_free(aud);
+            err = ff_cbs_alloc_unit_content(ctx, unit,
+                                            sizeof(H265RawAUD), NULL);
+            if (err < 0)
                 return err;
-            }
 
-            unit->content = aud;
+            err = cbs_h265_read_aud(ctx, &bc, unit->content);
+            if (err < 0)
+                return err;
         }
         break;
 
@@ -1272,7 +1242,7 @@ static int cbs_h2645_write_nal_unit(CodedBitstreamContext *ctx,
     unit->data_size = (put_bits_count(&pbc) + 7) / 8;
     flush_put_bits(&pbc);
 
-    err = av_reallocp(&unit->data, unit->data_size);
+    err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
     if (err < 0)
         return err;
 
@@ -1354,6 +1324,12 @@ static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
     if (err)
         return err;
 
+    frag->data_ref = av_buffer_create(data, dp, NULL, NULL, 0);
+    if (!frag->data_ref) {
+        av_freep(&data);
+        return AVERROR(ENOMEM);
+    }
+
     frag->data = data;
     frag->data_size = dp;
 
@@ -1402,7 +1378,6 @@ const CodedBitstreamType ff_cbs_type_h264 = {
     .write_unit        = &cbs_h2645_write_nal_unit,
     .assemble_fragment = &cbs_h2645_assemble_fragment,
 
-    .free_unit         = &cbs_h264_free_nal_unit,
     .close             = &cbs_h264_close,
 };
 
@@ -1416,6 +1391,5 @@ const CodedBitstreamType ff_cbs_type_h265 = {
     .write_unit        = &cbs_h2645_write_nal_unit,
     .assemble_fragment = &cbs_h2645_assemble_fragment,
 
-    .free_unit         = &cbs_h265_free_nal_unit,
     .close             = &cbs_h265_close,
 };
diff --git a/libavcodec/cbs_h265.h b/libavcodec/cbs_h265.h
index 6ee0725933..0628748f18 100644
--- a/libavcodec/cbs_h265.h
+++ b/libavcodec/cbs_h265.h
@@ -154,6 +154,7 @@ typedef struct H265RawVUI {
 typedef struct H265RawPSExtensionData {
     uint8_t *data;
     size_t bit_length;
+    AVBufferRef *data_ref;
 } H265RawPSExtensionData;
 
 typedef struct H265RawVPS {
@@ -512,6 +513,7 @@ typedef struct H265RawSlice {
     uint8_t *data;
     size_t   data_size;
     int      data_bit_start;
+    AVBufferRef *data_ref;
 } H265RawSlice;
 
 
diff --git a/libavcodec/cbs_internal.h b/libavcodec/cbs_internal.h
index 1a1c22f068..4c6f421d19 100644
--- a/libavcodec/cbs_internal.h
+++ b/libavcodec/cbs_internal.h
@@ -53,9 +53,6 @@ typedef struct CodedBitstreamType {
     int (*assemble_fragment)(CodedBitstreamContext *ctx,
                              CodedBitstreamFragment *frag);
 
-    // Free the content and data of a single unit.
-    void (*free_unit)(CodedBitstreamUnit *unit);
-
     // Free the codec internal state.
     void (*close)(CodedBitstreamContext *ctx);
 } CodedBitstreamType;
diff --git a/libavcodec/cbs_mpeg2.c b/libavcodec/cbs_mpeg2.c
index 5956f39335..3db10c5152 100644
--- a/libavcodec/cbs_mpeg2.c
+++ b/libavcodec/cbs_mpeg2.c
@@ -102,6 +102,21 @@
 #undef nextbits
 
 
+static void cbs_mpeg2_free_user_data(void *unit, uint8_t *content)
+{
+    MPEG2RawUserData *user = (MPEG2RawUserData*)content;
+    av_buffer_unref(&user->user_data_ref);
+    av_freep(&content);
+}
+
+static void cbs_mpeg2_free_slice(void *unit, uint8_t *content)
+{
+    MPEG2RawSlice *slice = (MPEG2RawSlice*)content;
+    av_buffer_unref(&slice->header.extra_information_ref);
+    av_buffer_unref(&slice->data_ref);
+    av_freep(&content);
+}
+
 static int cbs_mpeg2_split_fragment(CodedBitstreamContext *ctx,
                                     CodedBitstreamFragment *frag,
                                     int header)
@@ -138,7 +153,7 @@ static int cbs_mpeg2_split_fragment(CodedBitstreamContext *ctx,
         memset(unit_data + unit_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 
         err = ff_cbs_insert_unit_data(ctx, frag, i, unit_type,
-                                      unit_data, unit_size);
+                                      unit_data, unit_size, NULL);
         if (err < 0) {
             av_freep(&unit_data);
             return err;
@@ -168,25 +183,25 @@ static int cbs_mpeg2_read_unit(CodedBitstreamContext *ctx,
         MPEG2RawSlice *slice;
         int pos, len;
 
-        slice = av_mallocz(sizeof(*slice));
-        if (!slice)
-            return AVERROR(ENOMEM);
+        err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*slice),
+                                        &cbs_mpeg2_free_slice);
+        if (err < 0)
+            return err;
+        slice = unit->content;
+
         err = cbs_mpeg2_read_slice_header(ctx, &bc, &slice->header);
-        if (err < 0) {
-            av_free(slice);
+        if (err < 0)
             return err;
-        }
 
         pos = bitstream_tell(&bc);
         len = unit->data_size;
 
         slice->data_size = len - pos / 8;
-        slice->data = av_malloc(slice->data_size +
-                                AV_INPUT_BUFFER_PADDING_SIZE);
-        if (!slice->data) {
-            av_free(slice);
+        slice->data_ref  = av_buffer_alloc(slice->data_size +
+                                           AV_INPUT_BUFFER_PADDING_SIZE);
+        if (!slice->data_ref)
             return AVERROR(ENOMEM);
-        }
+        slice->data = slice->data_ref->data;
 
         memcpy(slice->data,
                unit->data + pos / 8, slice->data_size);
@@ -194,30 +209,29 @@ static int cbs_mpeg2_read_unit(CodedBitstreamContext *ctx,
                AV_INPUT_BUFFER_PADDING_SIZE);
         slice->data_bit_start = pos % 8;
 
-        unit->content = slice;
-
     } else {
         switch (unit->type) {
-#define START(start_code, type, func) \
+#define START(start_code, type, read_func, free_func) \
         case start_code: \
             { \
                 type *header; \
-                header = av_mallocz(sizeof(*header)); \
-                if (!header) \
-                    return AVERROR(ENOMEM); \
-                err = cbs_mpeg2_read_ ## func(ctx, &bc, header); \
-                if (err < 0) { \
-                    av_free(header); \
+                err = ff_cbs_alloc_unit_content(ctx, unit, \
+                                                sizeof(*header), free_func); \
+                if (err < 0) \
+                    return err; \
+                header = unit->content; \
+                err = cbs_mpeg2_read_ ## read_func(ctx, &bc, header); \
+                if (err < 0) \
                     return err; \
-                } \
-                unit->content = header; \
             } \
             break;
-            START(0x00, MPEG2RawPictureHeader,  picture_header);
-            START(0xb2, MPEG2RawUserData,       user_data);
-            START(0xb3, MPEG2RawSequenceHeader, sequence_header);
-            START(0xb5, MPEG2RawExtensionData,  extension_data);
-            START(0xb8, MPEG2RawGroupOfPicturesHeader, group_of_pictures_header);
+            START(0x00, MPEG2RawPictureHeader,  picture_header,  NULL);
+            START(0xb2, MPEG2RawUserData,       user_data,
+                                            &cbs_mpeg2_free_user_data);
+            START(0xb3, MPEG2RawSequenceHeader, sequence_header, NULL);
+            START(0xb5, MPEG2RawExtensionData,  extension_data,  NULL);
+            START(0xb8, MPEG2RawGroupOfPicturesHeader,
+                                       group_of_pictures_header, NULL);
 #undef START
         default:
             av_log(ctx->log_ctx, AV_LOG_ERROR, "Unknown start code %02"PRIx32".\n",
@@ -335,7 +349,7 @@ static int cbs_mpeg2_write_unit(CodedBitstreamContext *ctx,
     unit->data_size = (put_bits_count(&pbc) + 7) / 8;
     flush_put_bits(&pbc);
 
-    err = av_reallocp(&unit->data, unit->data_size);
+    err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
     if (err < 0)
         return err;
 
@@ -355,9 +369,10 @@ static int cbs_mpeg2_assemble_fragment(CodedBitstreamContext *ctx,
     for (i = 0; i < frag->nb_units; i++)
         size += 3 + frag->units[i].data_size;
 
-    data = av_malloc(size);
-    if (!data)
+    frag->data_ref = av_buffer_alloc(size);
+    if (!frag->data_ref)
         return AVERROR(ENOMEM);
+    data = frag->data_ref->data;
 
     dp = 0;
     for (i = 0; i < frag->nb_units; i++) {
@@ -379,19 +394,6 @@ static int cbs_mpeg2_assemble_fragment(CodedBitstreamContext *ctx,
     return 0;
 }
 
-static void cbs_mpeg2_free_unit(CodedBitstreamUnit *unit)
-{
-    if (MPEG2_START_IS_SLICE(unit->type)) {
-        MPEG2RawSlice *slice = unit->content;
-        av_freep(&slice->data);
-        av_freep(&slice->header.extra_information);
-    } else if (unit->type == MPEG2_START_USER_DATA) {
-        MPEG2RawUserData *user = unit->content;
-        av_freep(&user->user_data);
-    }
-    av_freep(&unit->content);
-}
-
 static void cbs_mpeg2_close(CodedBitstreamContext *ctx)
 {
     CodedBitstreamMPEG2Context *priv = ctx->priv_data;
@@ -409,6 +411,5 @@ const CodedBitstreamType ff_cbs_type_mpeg2 = {
     .write_unit        = &cbs_mpeg2_write_unit,
     .assemble_fragment = &cbs_mpeg2_assemble_fragment,
 
-    .free_unit         = &cbs_mpeg2_free_unit,
     .close             = &cbs_mpeg2_close,
 };
diff --git a/libavcodec/cbs_mpeg2.h b/libavcodec/cbs_mpeg2.h
index 0b6cb998dc..31e57ce36b 100644
--- a/libavcodec/cbs_mpeg2.h
+++ b/libavcodec/cbs_mpeg2.h
@@ -22,6 +22,8 @@
 #include <stddef.h>
 #include <stdint.h>
 
+#include "libavutil/buffer.h"
+
 
 enum {
     MPEG2_START_PICTURE         = 0x00,
@@ -76,6 +78,7 @@ typedef struct MPEG2RawUserData {
 
     uint8_t *user_data;
     size_t user_data_length;
+    AVBufferRef *user_data_ref;
 } MPEG2RawUserData;
 
 typedef struct MPEG2RawSequenceExtension {
@@ -195,6 +198,7 @@ typedef struct MPEG2RawSliceHeader {
 
     size_t extra_information_length;
     uint8_t *extra_information;
+    AVBufferRef *extra_information_ref;
 } MPEG2RawSliceHeader;
 
 typedef struct MPEG2RawSlice {
@@ -203,6 +207,7 @@ typedef struct MPEG2RawSlice {
     uint8_t *data;
     size_t   data_size;
     int      data_bit_start;
+    AVBufferRef *data_ref;
 } MPEG2RawSlice;
 
 
diff --git a/libavcodec/cbs_mpeg2_syntax_template.c b/libavcodec/cbs_mpeg2_syntax_template.c
index 4aa1eb3c06..b61fc8bc24 100644
--- a/libavcodec/cbs_mpeg2_syntax_template.c
+++ b/libavcodec/cbs_mpeg2_syntax_template.c
@@ -71,9 +71,10 @@ static int FUNC(user_data)(CodedBitstreamContext *ctx, RWContext *rw,
     av_assert0(k % 8 == 0);
     current->user_data_length = k /= 8;
     if (k > 0) {
-        current->user_data = av_malloc(k);
-        if (!current->user_data)
+        current->user_data_ref = av_buffer_alloc(k);
+        if (!current->user_data_ref)
             return AVERROR(ENOMEM);
+        current->user_data = current->user_data_ref->data;
     }
 #endif
 
diff --git a/libavcodec/h264_metadata_bsf.c b/libavcodec/h264_metadata_bsf.c
index 2b579e9d3d..6a1904e315 100644
--- a/libavcodec/h264_metadata_bsf.c
+++ b/libavcodec/h264_metadata_bsf.c
@@ -270,7 +270,7 @@ static int h264_metadata_filter(AVBSFContext *bsf, AVPacket *out)
             aud->primary_pic_type = j;
 
             err = ff_cbs_insert_unit_content(ctx->cbc, au,
-                                             0, H264_NAL_AUD, aud);
+                                             0, H264_NAL_AUD, aud, NULL);
             if (err < 0) {
                 av_log(bsf, AV_LOG_ERROR, "Failed to insert AUD.\n");
                 goto fail;
@@ -314,8 +314,8 @@ static int h264_metadata_filter(AVBSFContext *bsf, AVPacket *out)
 
             sei->nal_unit_header.nal_unit_type = H264_NAL_SEI;
 
-            err = ff_cbs_insert_unit_content(ctx->cbc, au,
-                                             sei_pos, H264_NAL_SEI, sei);
+            err = ff_cbs_insert_unit_content(ctx->cbc, au, sei_pos,
+                                             H264_NAL_SEI, sei, NULL);
             if (err < 0) {
                 av_log(bsf, AV_LOG_ERROR, "Failed to insert SEI.\n");
                 goto fail;
diff --git a/libavcodec/h265_metadata_bsf.c b/libavcodec/h265_metadata_bsf.c
index ffaf7f2873..cb73210c0d 100644
--- a/libavcodec/h265_metadata_bsf.c
+++ b/libavcodec/h265_metadata_bsf.c
@@ -289,7 +289,7 @@ static int h265_metadata_filter(AVBSFContext *bsf, AVPacket *out)
             aud->pic_type = pic_type;
 
             err = ff_cbs_insert_unit_content(ctx->cbc, au,
-                                             0, HEVC_NAL_AUD, aud);
+                                             0, HEVC_NAL_AUD, aud, NULL);
             if (err) {
                 av_log(bsf, AV_LOG_ERROR, "Failed to insert AUD.\n");
                 goto fail;
diff --git a/libavcodec/mpeg2_metadata_bsf.c b/libavcodec/mpeg2_metadata_bsf.c
index 4923932864..668d70ea2f 100644
--- a/libavcodec/mpeg2_metadata_bsf.c
+++ b/libavcodec/mpeg2_metadata_bsf.c
@@ -167,7 +167,8 @@ static int mpeg2_metadata_update_fragment(AVBSFContext *bsf,
 
         err = ff_cbs_insert_unit_content(ctx->cbc, frag, se_pos + 1,
                                          MPEG2_START_EXTENSION,
-                                         &ctx->sequence_display_extension);
+                                         &ctx->sequence_display_extension,
+                                         NULL);
         if (err < 0) {
             av_log(bsf, AV_LOG_ERROR, "Failed to insert new sequence "
                    "display extension.\n");
diff --git a/libavcodec/vaapi_encode_h264.c b/libavcodec/vaapi_encode_h264.c
index a9f8832e20..74a6417e0d 100644
--- a/libavcodec/vaapi_encode_h264.c
+++ b/libavcodec/vaapi_encode_h264.c
@@ -133,7 +133,7 @@ static int vaapi_encode_h264_add_nal(AVCodecContext *avctx,
     int err;
 
     err = ff_cbs_insert_unit_content(priv->cbc, au, -1,
-                                     header->nal_unit_type, nal_unit);
+                                     header->nal_unit_type, nal_unit, NULL);
     if (err < 0) {
         av_log(avctx, AV_LOG_ERROR, "Failed to add NAL unit: "
                "type = %d.\n", header->nal_unit_type);
diff --git a/libavcodec/vaapi_encode_h265.c b/libavcodec/vaapi_encode_h265.c
index 52ac4a6875..9080aa82da 100644
--- a/libavcodec/vaapi_encode_h265.c
+++ b/libavcodec/vaapi_encode_h265.c
@@ -105,7 +105,7 @@ static int vaapi_encode_h265_add_nal(AVCodecContext *avctx,
     int err;
 
     err = ff_cbs_insert_unit_content(priv->cbc, au, -1,
-                                     header->nal_unit_type, nal_unit);
+                                     header->nal_unit_type, nal_unit, NULL);
     if (err < 0) {
         av_log(avctx, AV_LOG_ERROR, "Failed to add NAL unit: "
                "type = %d.\n", header->nal_unit_type);
diff --git a/libavcodec/vaapi_encode_mpeg2.c b/libavcodec/vaapi_encode_mpeg2.c
index 39f1e4b5f2..df26ed4c89 100644
--- a/libavcodec/vaapi_encode_mpeg2.c
+++ b/libavcodec/vaapi_encode_mpeg2.c
@@ -92,7 +92,7 @@ static int vaapi_encode_mpeg2_add_header(AVCodecContext *avctx,
     VAAPIEncodeMPEG2Context *priv = ctx->priv_data;
     int err;
 
-    err = ff_cbs_insert_unit_content(priv->cbc, frag, -1, type, header);
+    err = ff_cbs_insert_unit_content(priv->cbc, frag, -1, type, header, NULL);
     if (err < 0) {
         av_log(avctx, AV_LOG_ERROR, "Failed to add header: "
                "type = %d.\n", type);