]> git.sesse.net Git - ffmpeg/blobdiff - libavcodec/cbs_h2645.c
avcodec/idcinvideo: Mark decoder as init-threadsafe
[ffmpeg] / libavcodec / cbs_h2645.c
index 5dd371153a418a1155058b949129196d0ea52054..434322492c41ad3a8f13dba5aacd87aef6aaad74 100644 (file)
@@ -233,6 +233,16 @@ static int cbs_write_se_golomb(CodedBitstreamContext *ctx, PutBitContext *pbc,
     return 0;
 }
 
+// payload_extension_present() - true if we are before the last 1-bit
+// in the payload structure, which must be in the last byte.
+static int cbs_h265_payload_extension_present(GetBitContext *gbc, uint32_t payload_size,
+                                              int cur_pos)
+{
+    int bits_left = payload_size * 8 - cur_pos;
+    return (bits_left > 0 &&
+            (bits_left > 7 || show_bits(gbc, bits_left) & MAX_UINT_BITS(bits_left - 1)));
+}
+
 #define HEADER(name) do { \
         ff_cbs_trace_header(ctx, name); \
     } while (0)
@@ -250,18 +260,18 @@ static int cbs_write_se_golomb(CodedBitstreamContext *ctx, PutBitContext *pbc,
 #define SUBSCRIPTS(subs, ...) (subs > 0 ? ((int[subs + 1]){ subs, __VA_ARGS__ }) : NULL)
 
 #define u(width, name, range_min, range_max) \
-        xu(width, name, current->name, range_min, range_max, 0)
+        xu(width, name, current->name, range_min, range_max, 0)
 #define ub(width, name) \
-        xu(width, name, current->name, 0, MAX_UINT_BITS(width), 0)
+        xu(width, name, current->name, 0, MAX_UINT_BITS(width), 0)
 #define flag(name) ub(1, name)
 #define ue(name, range_min, range_max) \
-        xue(name, current->name, range_min, range_max, 0)
+        xue(name, current->name, range_min, range_max, 0)
 #define i(width, name, range_min, range_max) \
-        xi(width, name, current->name, range_min, range_max, 0)
+        xi(width, name, current->name, range_min, range_max, 0)
 #define ib(width, name) \
-        xi(width, name, current->name, MIN_INT_BITS(width), MAX_INT_BITS(width), 0)
+        xi(width, name, current->name, MIN_INT_BITS(width), MAX_INT_BITS(width), 0)
 #define se(name, range_min, range_max) \
-        xse(name, current->name, range_min, range_max, 0)
+        xse(name, current->name, range_min, range_max, 0)
 
 #define us(width, name, range_min, range_max, subs, ...) \
         xu(width, name, current->name, range_min, range_max, subs, __VA_ARGS__)
@@ -280,7 +290,7 @@ static int cbs_write_se_golomb(CodedBitstreamContext *ctx, PutBitContext *pbc,
 
 #define fixed(width, name, value) do { \
         av_unused uint32_t fixed_value = value; \
-        xu(width, name, fixed_value, value, value, 0); \
+        xu(width, name, fixed_value, value, value, 0); \
     } while (0)
 
 
@@ -398,10 +408,11 @@ static int cbs_h2645_read_more_rbsp_data(GetBitContext *gbc)
 
 #define infer(name, value) do { \
         if (current->name != (value)) { \
-            av_log(ctx->log_ctx, AV_LOG_WARNING, "Warning: " \
+            av_log(ctx->log_ctx, AV_LOG_ERROR, \
                    "%s does not match inferred value: " \
                    "%"PRId64", but should be %"PRId64".\n", \
                    #name, (int64_t)current->name, (int64_t)(value)); \
+            return AVERROR_INVALIDDATA; \
         } \
     } while (0)
 
@@ -443,117 +454,6 @@ static int cbs_h2645_read_more_rbsp_data(GetBitContext *gbc)
 #undef allocate
 
 
-static void cbs_h264_free_pps(void *opaque, uint8_t *content)
-{
-    H264RawPPS *pps = (H264RawPPS*)content;
-    av_buffer_unref(&pps->slice_group_id_ref);
-    av_freep(&content);
-}
-
-static void cbs_h264_free_sei_payload(H264RawSEIPayload *payload)
-{
-    switch (payload->payload_type) {
-    case H264_SEI_TYPE_BUFFERING_PERIOD:
-    case H264_SEI_TYPE_PIC_TIMING:
-    case H264_SEI_TYPE_PAN_SCAN_RECT:
-    case H264_SEI_TYPE_RECOVERY_POINT:
-    case H264_SEI_TYPE_DISPLAY_ORIENTATION:
-    case H264_SEI_TYPE_MASTERING_DISPLAY_COLOUR_VOLUME:
-    case H264_SEI_TYPE_ALTERNATIVE_TRANSFER:
-        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_sei(void *opaque, uint8_t *content)
-{
-    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_slice(void *opaque, uint8_t *content)
-{
-    H264RawSlice *slice = (H264RawSlice*)content;
-    av_buffer_unref(&slice->data_ref);
-    av_freep(&content);
-}
-
-static void cbs_h265_free_vps(void *opaque, uint8_t *content)
-{
-    H265RawVPS *vps = (H265RawVPS*)content;
-    av_buffer_unref(&vps->extension_data.data_ref);
-    av_freep(&content);
-}
-
-static void cbs_h265_free_sps(void *opaque, 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 *opaque, 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 *opaque, uint8_t *content)
-{
-    H265RawSlice *slice = (H265RawSlice*)content;
-    av_buffer_unref(&slice->data_ref);
-    av_freep(&content);
-}
-
-static void cbs_h265_free_sei_payload(H265RawSEIPayload *payload)
-{
-    switch (payload->payload_type) {
-    case HEVC_SEI_TYPE_BUFFERING_PERIOD:
-    case HEVC_SEI_TYPE_PICTURE_TIMING:
-    case HEVC_SEI_TYPE_PAN_SCAN_RECT:
-    case HEVC_SEI_TYPE_RECOVERY_POINT:
-    case HEVC_SEI_TYPE_DISPLAY_ORIENTATION:
-    case HEVC_SEI_TYPE_ACTIVE_PARAMETER_SETS:
-    case HEVC_SEI_TYPE_DECODED_PICTURE_HASH:
-    case HEVC_SEI_TYPE_TIME_CODE:
-    case HEVC_SEI_TYPE_MASTERING_DISPLAY_INFO:
-    case HEVC_SEI_TYPE_CONTENT_LIGHT_LEVEL_INFO:
-    case HEVC_SEI_TYPE_ALTERNATIVE_TRANSFER_CHARACTERISTICS:
-    case HEVC_SEI_TYPE_ALPHA_CHANNEL_INFO:
-        break;
-    case HEVC_SEI_TYPE_USER_DATA_REGISTERED_ITU_T_T35:
-        av_buffer_unref(&payload->payload.user_data_registered.data_ref);
-        break;
-    case HEVC_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_h265_free_sei(void *opaque, uint8_t *content)
-{
-    H265RawSEI *sei = (H265RawSEI*)content;
-    int i;
-    for (i = 0; i < sei->payload_count; i++)
-        cbs_h265_free_sei_payload(&sei->payload[i]);
-    av_freep(&content);
-}
-
 static int cbs_h2645_fragment_add_nals(CodedBitstreamContext *ctx,
                                        CodedBitstreamFragment *frag,
                                        const H2645Packet *packet)
@@ -565,15 +465,21 @@ static int cbs_h2645_fragment_add_nals(CodedBitstreamContext *ctx,
         AVBufferRef *ref;
         size_t size = nal->size;
 
+        if (nal->nuh_layer_id > 0)
+            continue;
+
         // Remove trailing zeroes.
         while (size > 0 && nal->data[size - 1] == 0)
             --size;
-        av_assert0(size > 0);
+        if (size == 0) {
+            av_log(ctx->log_ctx, AV_LOG_VERBOSE, "Discarding empty 0 NAL unit\n");
+            continue;
+        }
 
         ref = (nal->data == nal->raw_data) ? frag->data_ref
                                            : packet->rbsp.rbsp_buffer_ref;
 
-        err = ff_cbs_insert_unit_data(ctx, frag, -1, nal->type,
+        err = ff_cbs_insert_unit_data(frag, -1, nal->type,
                             (uint8_t*)nal->data, size, ref);
         if (err < 0)
             return err;
@@ -610,7 +516,7 @@ static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
         version = bytestream2_get_byte(&gbc);
         if (version != 1) {
             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid AVCC header: "
-                   "first byte %u.", version);
+                   "first byte %u.\n", version);
             return AVERROR_INVALIDDATA;
         }
 
@@ -685,7 +591,7 @@ static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
         version = bytestream2_get_byte(&gbc);
         if (version != 1) {
             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid HVCC header: "
-                   "first byte %u.", version);
+                   "first byte %u.\n", version);
             return AVERROR_INVALIDDATA;
         }
 
@@ -748,23 +654,23 @@ static int cbs_h26 ## h26n ## _replace_ ## ps_var(CodedBitstreamContext *ctx, \
     CodedBitstreamH26 ## h26n ## Context *priv = ctx->priv_data; \
     H26 ## h26n ## Raw ## ps_name *ps_var = unit->content; \
     unsigned int id = ps_var->id_element; \
+    int err; \
     if (id >= FF_ARRAY_ELEMS(priv->ps_var)) { \
         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid " #ps_name \
                " id : %d.\n", id); \
         return AVERROR_INVALIDDATA; \
     } \
+    err = ff_cbs_make_unit_refcounted(ctx, unit); \
+    if (err < 0) \
+        return err; \
     if (priv->ps_var[id] == priv->active_ ## ps_var) \
         priv->active_ ## ps_var = NULL ; \
     av_buffer_unref(&priv->ps_var ## _ref[id]); \
-    if (unit->content_ref) \
-        priv->ps_var ## _ref[id] = av_buffer_ref(unit->content_ref); \
-    else \
-        priv->ps_var ## _ref[id] = av_buffer_alloc(sizeof(*ps_var)); \
+    av_assert0(unit->content_ref); \
+    priv->ps_var ## _ref[id] = av_buffer_ref(unit->content_ref); \
     if (!priv->ps_var ## _ref[id]) \
         return AVERROR(ENOMEM); \
     priv->ps_var[id] = (H26 ## h26n ## Raw ## ps_name *)priv->ps_var ## _ref[id]->data; \
-    if (!unit->content_ref) \
-        memcpy(priv->ps_var[id], ps_var, sizeof(*ps_var)); \
     return 0; \
 }
 
@@ -784,15 +690,14 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
     if (err < 0)
         return err;
 
+    err = ff_cbs_alloc_unit_content2(ctx, unit);
+    if (err < 0)
+        return err;
+
     switch (unit->type) {
     case H264_NAL_SPS:
         {
-            H264RawSPS *sps;
-
-            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*sps), NULL);
-            if (err < 0)
-                return err;
-            sps = unit->content;
+            H264RawSPS *sps = unit->content;
 
             err = cbs_h264_read_sps(ctx, &gbc, sps);
             if (err < 0)
@@ -806,12 +711,6 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
 
     case H264_NAL_SPS_EXT:
         {
-            err = ff_cbs_alloc_unit_content(ctx, unit,
-                                            sizeof(H264RawSPSExtension),
-                                            NULL);
-            if (err < 0)
-                return err;
-
             err = cbs_h264_read_sps_extension(ctx, &gbc, unit->content);
             if (err < 0)
                 return err;
@@ -820,13 +719,7 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
 
     case H264_NAL_PPS:
         {
-            H264RawPPS *pps;
-
-            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*pps),
-                                            &cbs_h264_free_pps);
-            if (err < 0)
-                return err;
-            pps = unit->content;
+            H264RawPPS *pps = unit->content;
 
             err = cbs_h264_read_pps(ctx, &gbc, pps);
             if (err < 0)
@@ -842,28 +735,18 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
     case H264_NAL_IDR_SLICE:
     case H264_NAL_AUXILIARY_SLICE:
         {
-            H264RawSlice *slice;
+            H264RawSlice *slice = unit->content;
             int pos, len;
 
-            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, &gbc, &slice->header);
             if (err < 0)
                 return err;
 
+            if (!cbs_h2645_read_more_rbsp_data(&gbc))
+                return AVERROR_INVALIDDATA;
+
             pos = get_bits_count(&gbc);
             len = unit->data_size;
-            if (!unit->data[len - 1]) {
-                int z;
-                for (z = 0; z < len && !unit->data[len - z - 1]; z++);
-                av_log(ctx->log_ctx, AV_LOG_DEBUG, "Deleted %d trailing zeroes "
-                       "from slice data.\n", z);
-                len -= z;
-            }
 
             slice->data_size = len - pos / 8;
             slice->data_ref  = av_buffer_ref(unit->data_ref);
@@ -876,11 +759,6 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
 
     case H264_NAL_AUD:
         {
-            err = ff_cbs_alloc_unit_content(ctx, unit,
-                                            sizeof(H264RawAUD), NULL);
-            if (err < 0)
-                return err;
-
             err = cbs_h264_read_aud(ctx, &gbc, unit->content);
             if (err < 0)
                 return err;
@@ -889,11 +767,6 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
 
     case H264_NAL_SEI:
         {
-            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(H264RawSEI),
-                                            &cbs_h264_free_sei);
-            if (err < 0)
-                return err;
-
             err = cbs_h264_read_sei(ctx, &gbc, unit->content);
             if (err < 0)
                 return err;
@@ -902,11 +775,6 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
 
     case H264_NAL_FILLER_DATA:
         {
-            err = ff_cbs_alloc_unit_content(ctx, unit,
-                                            sizeof(H264RawFiller), NULL);
-            if (err < 0)
-                return err;
-
             err = cbs_h264_read_filler(ctx, &gbc, unit->content);
             if (err < 0)
                 return err;
@@ -916,12 +784,6 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
     case H264_NAL_END_SEQUENCE:
     case H264_NAL_END_STREAM:
         {
-            err = ff_cbs_alloc_unit_content(ctx, unit,
-                                            sizeof(H264RawNALUnitHeader),
-                                            NULL);
-            if (err < 0)
-                return err;
-
             err = (unit->type == H264_NAL_END_SEQUENCE ?
                    cbs_h264_read_end_of_sequence :
                    cbs_h264_read_end_of_stream)(ctx, &gbc, unit->content);
@@ -947,16 +809,14 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
     if (err < 0)
         return err;
 
+    err = ff_cbs_alloc_unit_content2(ctx, unit);
+    if (err < 0)
+        return err;
+
     switch (unit->type) {
     case HEVC_NAL_VPS:
         {
-            H265RawVPS *vps;
-
-            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*vps),
-                                            &cbs_h265_free_vps);
-            if (err < 0)
-                return err;
-            vps = unit->content;
+            H265RawVPS *vps = unit->content;
 
             err = cbs_h265_read_vps(ctx, &gbc, vps);
             if (err < 0)
@@ -969,13 +829,7 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
         break;
     case HEVC_NAL_SPS:
         {
-            H265RawSPS *sps;
-
-            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*sps),
-                                            &cbs_h265_free_sps);
-            if (err < 0)
-                return err;
-            sps = unit->content;
+            H265RawSPS *sps = unit->content;
 
             err = cbs_h265_read_sps(ctx, &gbc, sps);
             if (err < 0)
@@ -989,13 +843,7 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
 
     case HEVC_NAL_PPS:
         {
-            H265RawPPS *pps;
-
-            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(*pps),
-                                            &cbs_h265_free_pps);
-            if (err < 0)
-                return err;
-            pps = unit->content;
+            H265RawPPS *pps = unit->content;
 
             err = cbs_h265_read_pps(ctx, &gbc, pps);
             if (err < 0)
@@ -1024,28 +872,18 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
     case HEVC_NAL_IDR_N_LP:
     case HEVC_NAL_CRA_NUT:
         {
-            H265RawSlice *slice;
+            H265RawSlice *slice = unit->content;
             int pos, len;
 
-            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, &gbc, &slice->header);
             if (err < 0)
                 return err;
 
+            if (!cbs_h2645_read_more_rbsp_data(&gbc))
+                return AVERROR_INVALIDDATA;
+
             pos = get_bits_count(&gbc);
             len = unit->data_size;
-            if (!unit->data[len - 1]) {
-                int z;
-                for (z = 0; z < len && !unit->data[len - z - 1]; z++);
-                av_log(ctx->log_ctx, AV_LOG_DEBUG, "Deleted %d trailing zeroes "
-                       "from slice data.\n", z);
-                len -= z;
-            }
 
             slice->data_size = len - pos / 8;
             slice->data_ref  = av_buffer_ref(unit->data_ref);
@@ -1058,11 +896,6 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
 
     case HEVC_NAL_AUD:
         {
-            err = ff_cbs_alloc_unit_content(ctx, unit,
-                                            sizeof(H265RawAUD), NULL);
-            if (err < 0)
-                return err;
-
             err = cbs_h265_read_aud(ctx, &gbc, unit->content);
             if (err < 0)
                 return err;
@@ -1072,12 +905,6 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
     case HEVC_NAL_SEI_PREFIX:
     case HEVC_NAL_SEI_SUFFIX:
         {
-            err = ff_cbs_alloc_unit_content(ctx, unit, sizeof(H265RawSEI),
-                                            &cbs_h265_free_sei);
-
-            if (err < 0)
-                return err;
-
             err = cbs_h265_read_sei(ctx, &gbc, unit->content,
                                     unit->type == HEVC_NAL_SEI_PREFIX);
 
@@ -1101,7 +928,7 @@ static int cbs_h2645_write_slice_data(CodedBitstreamContext *ctx,
     const uint8_t *pos = data + data_bit_start / 8;
 
     av_assert0(data_bit_start >= 0 &&
-               8 * data_size > data_bit_start);
+               data_size > data_bit_start / 8);
 
     if (data_size * 8 + 8 > put_bits_left(pbc))
         return AVERROR(ENOSPC);
@@ -1380,65 +1207,6 @@ static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
     return 0;
 }
 
-static int cbs_h2645_write_nal_unit(CodedBitstreamContext *ctx,
-                                    CodedBitstreamUnit *unit)
-{
-    CodedBitstreamH2645Context *priv = ctx->priv_data;
-    enum AVCodecID codec_id = ctx->codec->codec_id;
-    PutBitContext pbc;
-    int err;
-
-    if (!priv->write_buffer) {
-        // Initial write buffer size is 1MB.
-        priv->write_buffer_size = 1024 * 1024;
-
-    reallocate_and_try_again:
-        err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
-        if (err < 0) {
-            av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
-                   "sufficiently large write buffer (last attempt "
-                   "%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
-            return err;
-        }
-    }
-
-    init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
-
-    if (codec_id == AV_CODEC_ID_H264)
-        err = cbs_h264_write_nal_unit(ctx, unit, &pbc);
-    else
-        err = cbs_h265_write_nal_unit(ctx, unit, &pbc);
-
-    if (err == AVERROR(ENOSPC)) {
-        // Overflow.
-        priv->write_buffer_size *= 2;
-        goto reallocate_and_try_again;
-    }
-    // Overflow but we didn't notice.
-    av_assert0(put_bits_count(&pbc) <= 8 * priv->write_buffer_size);
-
-    if (err < 0) {
-        // Write failed for some other reason.
-        return err;
-    }
-
-    if (put_bits_count(&pbc) % 8)
-        unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
-    else
-        unit->data_bit_padding = 0;
-
-    unit->data_size = (put_bits_count(&pbc) + 7) / 8;
-    flush_put_bits(&pbc);
-
-    err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
-    if (err < 0)
-        return err;
-
-    memcpy(unit->data, priv->write_buffer, unit->data_size);
-
-    return 0;
-}
-
 static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
                                        CodedBitstreamFragment *frag)
 {
@@ -1454,7 +1222,7 @@ static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
     max_size = 0;
     for (i = 0; i < frag->nb_units; i++) {
         // Start code + content with worst-case emulation prevention.
-        max_size += 3 + frag->units[i].data_size * 3 / 2;
+        max_size += 4 + frag->units[i].data_size * 3 / 2;
     }
 
     data = av_realloc(NULL, max_size + AV_INPUT_BUFFER_PADDING_SIZE);
@@ -1526,6 +1294,24 @@ static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
     return 0;
 }
 
+static void cbs_h264_flush(CodedBitstreamContext *ctx)
+{
+    CodedBitstreamH264Context *h264 = ctx->priv_data;
+
+    for (int i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++) {
+        av_buffer_unref(&h264->sps_ref[i]);
+        h264->sps[i] = NULL;
+    }
+    for (int i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++) {
+        av_buffer_unref(&h264->pps_ref[i]);
+        h264->pps[i] = NULL;
+    }
+
+    h264->active_sps = NULL;
+    h264->active_pps = NULL;
+    h264->last_slice_nal_unit_type = 0;
+}
+
 static void cbs_h264_close(CodedBitstreamContext *ctx)
 {
     CodedBitstreamH264Context *h264 = ctx->priv_data;
@@ -1533,14 +1319,34 @@ static void cbs_h264_close(CodedBitstreamContext *ctx)
 
     ff_h2645_packet_uninit(&h264->common.read_packet);
 
-    av_freep(&h264->common.write_buffer);
-
     for (i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++)
         av_buffer_unref(&h264->sps_ref[i]);
     for (i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++)
         av_buffer_unref(&h264->pps_ref[i]);
 }
 
+static void cbs_h265_flush(CodedBitstreamContext *ctx)
+{
+    CodedBitstreamH265Context *h265 = ctx->priv_data;
+
+    for (int i = 0; i < FF_ARRAY_ELEMS(h265->vps); i++) {
+        av_buffer_unref(&h265->vps_ref[i]);
+        h265->vps[i] = NULL;
+    }
+    for (int i = 0; i < FF_ARRAY_ELEMS(h265->sps); i++) {
+        av_buffer_unref(&h265->sps_ref[i]);
+        h265->sps[i] = NULL;
+    }
+    for (int i = 0; i < FF_ARRAY_ELEMS(h265->pps); i++) {
+        av_buffer_unref(&h265->pps_ref[i]);
+        h265->pps[i] = NULL;
+    }
+
+    h265->active_vps = NULL;
+    h265->active_sps = NULL;
+    h265->active_pps = NULL;
+}
+
 static void cbs_h265_close(CodedBitstreamContext *ctx)
 {
     CodedBitstreamH265Context *h265 = ctx->priv_data;
@@ -1548,8 +1354,6 @@ static void cbs_h265_close(CodedBitstreamContext *ctx)
 
     ff_h2645_packet_uninit(&h265->common.read_packet);
 
-    av_freep(&h265->common.write_buffer);
-
     for (i = 0; i < FF_ARRAY_ELEMS(h265->vps); i++)
         av_buffer_unref(&h265->vps_ref[i]);
     for (i = 0; i < FF_ARRAY_ELEMS(h265->sps); i++)
@@ -1558,16 +1362,163 @@ static void cbs_h265_close(CodedBitstreamContext *ctx)
         av_buffer_unref(&h265->pps_ref[i]);
 }
 
+static void cbs_h264_free_sei_payload(H264RawSEIPayload *payload)
+{
+    switch (payload->payload_type) {
+    case H264_SEI_TYPE_BUFFERING_PERIOD:
+    case H264_SEI_TYPE_PIC_TIMING:
+    case H264_SEI_TYPE_PAN_SCAN_RECT:
+    case H264_SEI_TYPE_RECOVERY_POINT:
+    case H264_SEI_TYPE_DISPLAY_ORIENTATION:
+    case H264_SEI_TYPE_MASTERING_DISPLAY_COLOUR_VOLUME:
+    case H264_SEI_TYPE_ALTERNATIVE_TRANSFER:
+        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_sei(void *opaque, uint8_t *content)
+{
+    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 const CodedBitstreamUnitTypeDescriptor cbs_h264_unit_types[] = {
+    CBS_UNIT_TYPE_POD(H264_NAL_SPS,     H264RawSPS),
+    CBS_UNIT_TYPE_POD(H264_NAL_SPS_EXT, H264RawSPSExtension),
+
+    CBS_UNIT_TYPE_INTERNAL_REF(H264_NAL_PPS, H264RawPPS, slice_group_id),
+
+    {
+        .nb_unit_types  = 3,
+        .unit_types     = {
+            H264_NAL_IDR_SLICE,
+            H264_NAL_SLICE,
+            H264_NAL_AUXILIARY_SLICE,
+        },
+        .content_type   = CBS_CONTENT_TYPE_INTERNAL_REFS,
+        .content_size   = sizeof(H264RawSlice),
+        .nb_ref_offsets = 1,
+        .ref_offsets    = { offsetof(H264RawSlice, data) },
+    },
+
+    CBS_UNIT_TYPE_POD(H264_NAL_AUD,          H264RawAUD),
+    CBS_UNIT_TYPE_POD(H264_NAL_FILLER_DATA,  H264RawFiller),
+    CBS_UNIT_TYPE_POD(H264_NAL_END_SEQUENCE, H264RawNALUnitHeader),
+    CBS_UNIT_TYPE_POD(H264_NAL_END_STREAM,   H264RawNALUnitHeader),
+
+    CBS_UNIT_TYPE_COMPLEX(H264_NAL_SEI, H264RawSEI, &cbs_h264_free_sei),
+
+    CBS_UNIT_TYPE_END_OF_LIST
+};
+
+static void cbs_h265_free_sei_payload(H265RawSEIPayload *payload)
+{
+    switch (payload->payload_type) {
+    case HEVC_SEI_TYPE_BUFFERING_PERIOD:
+    case HEVC_SEI_TYPE_PICTURE_TIMING:
+    case HEVC_SEI_TYPE_PAN_SCAN_RECT:
+    case HEVC_SEI_TYPE_RECOVERY_POINT:
+    case HEVC_SEI_TYPE_DISPLAY_ORIENTATION:
+    case HEVC_SEI_TYPE_ACTIVE_PARAMETER_SETS:
+    case HEVC_SEI_TYPE_DECODED_PICTURE_HASH:
+    case HEVC_SEI_TYPE_TIME_CODE:
+    case HEVC_SEI_TYPE_MASTERING_DISPLAY_INFO:
+    case HEVC_SEI_TYPE_CONTENT_LIGHT_LEVEL_INFO:
+    case HEVC_SEI_TYPE_ALTERNATIVE_TRANSFER_CHARACTERISTICS:
+    case HEVC_SEI_TYPE_ALPHA_CHANNEL_INFO:
+        break;
+    case HEVC_SEI_TYPE_USER_DATA_REGISTERED_ITU_T_T35:
+        av_buffer_unref(&payload->payload.user_data_registered.data_ref);
+        break;
+    case HEVC_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;
+    }
+    av_buffer_unref(&payload->extension_data.data_ref);
+}
+
+static void cbs_h265_free_sei(void *opaque, uint8_t *content)
+{
+    H265RawSEI *sei = (H265RawSEI*)content;
+    int i;
+    for (i = 0; i < sei->payload_count; i++)
+        cbs_h265_free_sei_payload(&sei->payload[i]);
+    av_freep(&content);
+}
+
+static const CodedBitstreamUnitTypeDescriptor cbs_h265_unit_types[] = {
+    CBS_UNIT_TYPE_INTERNAL_REF(HEVC_NAL_VPS, H265RawVPS, extension_data.data),
+    CBS_UNIT_TYPE_INTERNAL_REF(HEVC_NAL_SPS, H265RawSPS, extension_data.data),
+    CBS_UNIT_TYPE_INTERNAL_REF(HEVC_NAL_PPS, H265RawPPS, extension_data.data),
+
+    CBS_UNIT_TYPE_POD(HEVC_NAL_AUD, H265RawAUD),
+
+    {
+        // Slices of non-IRAP pictures.
+        .nb_unit_types         = CBS_UNIT_TYPE_RANGE,
+        .unit_type_range_start = HEVC_NAL_TRAIL_N,
+        .unit_type_range_end   = HEVC_NAL_RASL_R,
+
+        .content_type   = CBS_CONTENT_TYPE_INTERNAL_REFS,
+        .content_size   = sizeof(H265RawSlice),
+        .nb_ref_offsets = 1,
+        .ref_offsets    = { offsetof(H265RawSlice, data) },
+    },
+
+    {
+        // Slices of IRAP pictures.
+        .nb_unit_types         = CBS_UNIT_TYPE_RANGE,
+        .unit_type_range_start = HEVC_NAL_BLA_W_LP,
+        .unit_type_range_end   = HEVC_NAL_CRA_NUT,
+
+        .content_type   = CBS_CONTENT_TYPE_INTERNAL_REFS,
+        .content_size   = sizeof(H265RawSlice),
+        .nb_ref_offsets = 1,
+        .ref_offsets    = { offsetof(H265RawSlice, data) },
+    },
+
+    {
+        .nb_unit_types  = 2,
+        .unit_types     = {
+            HEVC_NAL_SEI_PREFIX,
+            HEVC_NAL_SEI_SUFFIX
+        },
+        .content_type   = CBS_CONTENT_TYPE_COMPLEX,
+        .content_size   = sizeof(H265RawSEI),
+        .content_free   = &cbs_h265_free_sei,
+    },
+
+    CBS_UNIT_TYPE_END_OF_LIST
+};
+
 const CodedBitstreamType ff_cbs_type_h264 = {
     .codec_id          = AV_CODEC_ID_H264,
 
     .priv_data_size    = sizeof(CodedBitstreamH264Context),
 
+    .unit_types        = cbs_h264_unit_types,
+
     .split_fragment    = &cbs_h2645_split_fragment,
     .read_unit         = &cbs_h264_read_nal_unit,
-    .write_unit        = &cbs_h2645_write_nal_unit,
+    .write_unit        = &cbs_h264_write_nal_unit,
     .assemble_fragment = &cbs_h2645_assemble_fragment,
 
+    .flush             = &cbs_h264_flush,
     .close             = &cbs_h264_close,
 };
 
@@ -1576,16 +1527,18 @@ const CodedBitstreamType ff_cbs_type_h265 = {
 
     .priv_data_size    = sizeof(CodedBitstreamH265Context),
 
+    .unit_types        = cbs_h265_unit_types,
+
     .split_fragment    = &cbs_h2645_split_fragment,
     .read_unit         = &cbs_h265_read_nal_unit,
-    .write_unit        = &cbs_h2645_write_nal_unit,
+    .write_unit        = &cbs_h265_write_nal_unit,
     .assemble_fragment = &cbs_h2645_assemble_fragment,
 
+    .flush             = &cbs_h265_flush,
     .close             = &cbs_h265_close,
 };
 
-int ff_cbs_h264_add_sei_message(CodedBitstreamContext *ctx,
-                                CodedBitstreamFragment *au,
+int ff_cbs_h264_add_sei_message(CodedBitstreamFragment *au,
                                 H264RawSEIPayload *payload)
 {
     H264RawSEI *sei = NULL;
@@ -1630,7 +1583,7 @@ int ff_cbs_h264_add_sei_message(CodedBitstreamContext *ctx,
                 break;
         }
 
-        err = ff_cbs_insert_unit_content(ctx, au, i, H264_NAL_SEI,
+        err = ff_cbs_insert_unit_content(au, i, H264_NAL_SEI,
                                          sei, sei_ref);
         av_buffer_unref(&sei_ref);
         if (err < 0)
@@ -1646,8 +1599,7 @@ fail:
     return err;
 }
 
-void ff_cbs_h264_delete_sei_message(CodedBitstreamContext *ctx,
-                                    CodedBitstreamFragment *au,
+void ff_cbs_h264_delete_sei_message(CodedBitstreamFragment *au,
                                     CodedBitstreamUnit *nal,
                                     int position)
 {
@@ -1665,7 +1617,7 @@ void ff_cbs_h264_delete_sei_message(CodedBitstreamContext *ctx,
                 break;
         }
 
-        ff_cbs_delete_unit(ctx, au, i);
+        ff_cbs_delete_unit(au, i);
     } else {
         cbs_h264_free_sei_payload(&sei->payload[position]);