git.sesse.net Git - ffmpeg/blob - libavcodec/cbs_h2645.c

   1 /*
   2  * This file is part of FFmpeg.
   3  *
   4  * FFmpeg is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Lesser General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2.1 of the License, or (at your option) any later version.
   8  *
   9  * FFmpeg is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * Lesser General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Lesser General Public
  15  * License along with FFmpeg; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  17  */
  18
  19 #include "libavutil/attributes.h"
  20 #include "libavutil/avassert.h"
  21
  22 #include "bytestream.h"
  23 #include "cbs.h"
  24 #include "cbs_internal.h"
  25 #include "cbs_h264.h"
  26 #include "cbs_h265.h"
  27 #include "h264.h"
  28 #include "h264_sei.h"
  29 #include "h2645_parse.h"
  30 #include "hevc.h"
  31 #include "hevc_sei.h"
  32
  33
  34 static int cbs_read_ue_golomb(CodedBitstreamContext *ctx, GetBitContext *gbc,
  35                               const char *name, const int *subscripts,
  36                               uint32_t *write_to,
  37                               uint32_t range_min, uint32_t range_max)
  38 {
  39     uint32_t value;
  40     int position, i, j;
  41     unsigned int k;
  42     char bits[65];
  43
  44     position = get_bits_count(gbc);
  45
  46     for (i = 0; i < 32; i++) {
  47         if (get_bits_left(gbc) < i + 1) {
  48             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid ue-golomb code at "
  49                    "%s: bitstream ended.\n", name);
  50             return AVERROR_INVALIDDATA;
  51         }
  52         k = get_bits1(gbc);
  53         bits[i] = k ? '1' : '0';
  54         if (k)
  55             break;
  56     }
  57     if (i >= 32) {
  58         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid ue-golomb code at "
  59                "%s: more than 31 zeroes.\n", name);
  60         return AVERROR_INVALIDDATA;
  61     }
  62     value = 1;
  63     for (j = 0; j < i; j++) {
  64         k = get_bits1(gbc);
  65         bits[i + j + 1] = k ? '1' : '0';
  66         value = value << 1 | k;
  67     }
  68     bits[i + j + 1] = 0;
  69     --value;
  70
  71     if (ctx->trace_enable)
  72         ff_cbs_trace_syntax_element(ctx, position, name, subscripts,
  73                                     bits, value);
  74
  75     if (value < range_min || value > range_max) {
  76         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
  77                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
  78                name, value, range_min, range_max);
  79         return AVERROR_INVALIDDATA;
  80     }
  81
  82     *write_to = value;
  83     return 0;
  84 }
  85
  86 static int cbs_read_se_golomb(CodedBitstreamContext *ctx, GetBitContext *gbc,
  87                               const char *name, const int *subscripts,
  88                               int32_t *write_to,
  89                               int32_t range_min, int32_t range_max)
  90 {
  91     int32_t value;
  92     int position, i, j;
  93     unsigned int k;
  94     uint32_t v;
  95     char bits[65];
  96
  97     position = get_bits_count(gbc);
  98
  99     for (i = 0; i < 32; i++) {
 100         if (get_bits_left(gbc) < i + 1) {
 101             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid se-golomb code at "
 102                    "%s: bitstream ended.\n", name);
 103             return AVERROR_INVALIDDATA;
 104         }
 105         k = get_bits1(gbc);
 106         bits[i] = k ? '1' : '0';
 107         if (k)
 108             break;
 109     }
 110     if (i >= 32) {
 111         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid se-golomb code at "
 112                "%s: more than 31 zeroes.\n", name);
 113         return AVERROR_INVALIDDATA;
 114     }
 115     v = 1;
 116     for (j = 0; j < i; j++) {
 117         k = get_bits1(gbc);
 118         bits[i + j + 1] = k ? '1' : '0';
 119         v = v << 1 | k;
 120     }
 121     bits[i + j + 1] = 0;
 122     if (v & 1)
 123         value = -(int32_t)(v / 2);
 124     else
 125         value = v / 2;
 126
 127     if (ctx->trace_enable)
 128         ff_cbs_trace_syntax_element(ctx, position, name, subscripts,
 129                                     bits, value);
 130
 131     if (value < range_min || value > range_max) {
 132         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 133                "%"PRId32", but must be in [%"PRId32",%"PRId32"].\n",
 134                name, value, range_min, range_max);
 135         return AVERROR_INVALIDDATA;
 136     }
 137
 138     *write_to = value;
 139     return 0;
 140 }
 141
 142 static int cbs_write_ue_golomb(CodedBitstreamContext *ctx, PutBitContext *pbc,
 143                                const char *name, const int *subscripts,
 144                                uint32_t value,
 145                                uint32_t range_min, uint32_t range_max)
 146 {
 147     int len;
 148
 149     if (value < range_min || value > range_max) {
 150         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 151                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
 152                name, value, range_min, range_max);
 153         return AVERROR_INVALIDDATA;
 154     }
 155     av_assert0(value != UINT32_MAX);
 156
 157     len = av_log2(value + 1);
 158     if (put_bits_left(pbc) < 2 * len + 1)
 159         return AVERROR(ENOSPC);
 160
 161     if (ctx->trace_enable) {
 162         char bits[65];
 163         int i;
 164
 165         for (i = 0; i < len; i++)
 166             bits[i] = '0';
 167         bits[len] = '1';
 168         for (i = 0; i < len; i++)
 169             bits[len + i + 1] = (value + 1) >> (len - i - 1) & 1 ? '1' : '0';
 170         bits[len + len + 1] = 0;
 171
 172         ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc),
 173                                     name, subscripts, bits, value);
 174     }
 175
 176     put_bits(pbc, len, 0);
 177     if (len + 1 < 32)
 178         put_bits(pbc, len + 1, value + 1);
 179     else
 180         put_bits32(pbc, value + 1);
 181
 182     return 0;
 183 }
 184
 185 static int cbs_write_se_golomb(CodedBitstreamContext *ctx, PutBitContext *pbc,
 186                                const char *name, const int *subscripts,
 187                                int32_t value,
 188                                int32_t range_min, int32_t range_max)
 189 {
 190     int len;
 191     uint32_t uvalue;
 192
 193     if (value < range_min || value > range_max) {
 194         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
 195                "%"PRId32", but must be in [%"PRId32",%"PRId32"].\n",
 196                name, value, range_min, range_max);
 197         return AVERROR_INVALIDDATA;
 198     }
 199     av_assert0(value != INT32_MIN);
 200
 201     if (value == 0)
 202         uvalue = 0;
 203     else if (value > 0)
 204         uvalue = 2 * (uint32_t)value - 1;
 205     else
 206         uvalue = 2 * (uint32_t)-value;
 207
 208     len = av_log2(uvalue + 1);
 209     if (put_bits_left(pbc) < 2 * len + 1)
 210         return AVERROR(ENOSPC);
 211
 212     if (ctx->trace_enable) {
 213         char bits[65];
 214         int i;
 215
 216         for (i = 0; i < len; i++)
 217             bits[i] = '0';
 218         bits[len] = '1';
 219         for (i = 0; i < len; i++)
 220             bits[len + i + 1] = (uvalue + 1) >> (len - i - 1) & 1 ? '1' : '0';
 221         bits[len + len + 1] = 0;
 222
 223         ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc),
 224                                     name, subscripts, bits, value);
 225     }
 226
 227     put_bits(pbc, len, 0);
 228     if (len + 1 < 32)
 229         put_bits(pbc, len + 1, uvalue + 1);
 230     else
 231         put_bits32(pbc, uvalue + 1);
 232
 233     return 0;
 234 }
 235
 236 // payload_extension_present() - true if we are before the last 1-bit
 237 // in the payload structure, which must be in the last byte.
 238 static int cbs_h265_payload_extension_present(GetBitContext *gbc, uint32_t payload_size,
 239                                               int cur_pos)
 240 {
 241     int bits_left = payload_size * 8 - cur_pos;
 242     return (bits_left > 0 &&
 243             (bits_left > 7 || show_bits(gbc, bits_left) & MAX_UINT_BITS(bits_left - 1)));
 244 }
 245
 246 #define HEADER(name) do { \
 247         ff_cbs_trace_header(ctx, name); \
 248     } while (0)
 249
 250 #define CHECK(call) do { \
 251         err = (call); \
 252         if (err < 0) \
 253             return err; \
 254     } while (0)
 255
 256 #define FUNC_NAME(rw, codec, name) cbs_ ## codec ## _ ## rw ## _ ## name
 257 #define FUNC_H264(rw, name) FUNC_NAME(rw, h264, name)
 258 #define FUNC_H265(rw, name) FUNC_NAME(rw, h265, name)
 259
 260 #define SUBSCRIPTS(subs, ...) (subs > 0 ? ((int[subs + 1]){ subs, __VA_ARGS__ }) : NULL)
 261
 262 #define u(width, name, range_min, range_max) \
 263         xu(width, name, current->name, range_min, range_max, 0, )
 264 #define ub(width, name) \
 265         xu(width, name, current->name, 0, MAX_UINT_BITS(width), 0, )
 266 #define flag(name) ub(1, name)
 267 #define ue(name, range_min, range_max) \
 268         xue(name, current->name, range_min, range_max, 0, )
 269 #define i(width, name, range_min, range_max) \
 270         xi(width, name, current->name, range_min, range_max, 0, )
 271 #define ib(width, name) \
 272         xi(width, name, current->name, MIN_INT_BITS(width), MAX_INT_BITS(width), 0, )
 273 #define se(name, range_min, range_max) \
 274         xse(name, current->name, range_min, range_max, 0, )
 275
 276 #define us(width, name, range_min, range_max, subs, ...) \
 277         xu(width, name, current->name, range_min, range_max, subs, __VA_ARGS__)
 278 #define ubs(width, name, subs, ...) \
 279         xu(width, name, current->name, 0, MAX_UINT_BITS(width), subs, __VA_ARGS__)
 280 #define flags(name, subs, ...) \
 281         xu(1, name, current->name, 0, 1, subs, __VA_ARGS__)
 282 #define ues(name, range_min, range_max, subs, ...) \
 283         xue(name, current->name, range_min, range_max, subs, __VA_ARGS__)
 284 #define is(width, name, range_min, range_max, subs, ...) \
 285         xi(width, name, current->name, range_min, range_max, subs, __VA_ARGS__)
 286 #define ibs(width, name, subs, ...) \
 287         xi(width, name, current->name, MIN_INT_BITS(width), MAX_INT_BITS(width), subs, __VA_ARGS__)
 288 #define ses(name, range_min, range_max, subs, ...) \
 289         xse(name, current->name, range_min, range_max, subs, __VA_ARGS__)
 290
 291 #define fixed(width, name, value) do { \
 292         av_unused uint32_t fixed_value = value; \
 293         xu(width, name, fixed_value, value, value, 0, ); \
 294     } while (0)
 295
 296
 297 #define READ
 298 #define READWRITE read
 299 #define RWContext GetBitContext
 300
 301 #define xu(width, name, var, range_min, range_max, subs, ...) do { \
 302         uint32_t value; \
 303         CHECK(ff_cbs_read_unsigned(ctx, rw, width, #name, \
 304                                    SUBSCRIPTS(subs, __VA_ARGS__), \
 305                                    &value, range_min, range_max)); \
 306         var = value; \
 307     } while (0)
 308 #define xue(name, var, range_min, range_max, subs, ...) do { \
 309         uint32_t value; \
 310         CHECK(cbs_read_ue_golomb(ctx, rw, #name, \
 311                                  SUBSCRIPTS(subs, __VA_ARGS__), \
 312                                  &value, range_min, range_max)); \
 313         var = value; \
 314     } while (0)
 315 #define xi(width, name, var, range_min, range_max, subs, ...) do { \
 316         int32_t value; \
 317         CHECK(ff_cbs_read_signed(ctx, rw, width, #name, \
 318                                  SUBSCRIPTS(subs, __VA_ARGS__), \
 319                                  &value, range_min, range_max)); \
 320         var = value; \
 321     } while (0)
 322 #define xse(name, var, range_min, range_max, subs, ...) do { \
 323         int32_t value; \
 324         CHECK(cbs_read_se_golomb(ctx, rw, #name, \
 325                                  SUBSCRIPTS(subs, __VA_ARGS__), \
 326                                  &value, range_min, range_max)); \
 327         var = value; \
 328     } while (0)
 329
 330
 331 #define infer(name, value) do { \
 332         current->name = value; \
 333     } while (0)
 334
 335 static int cbs_h2645_read_more_rbsp_data(GetBitContext *gbc)
 336 {
 337     int bits_left = get_bits_left(gbc);
 338     if (bits_left > 8)
 339         return 1;
 340     if (bits_left == 0)
 341         return 0;
 342     if (show_bits(gbc, bits_left) & MAX_UINT_BITS(bits_left - 1))
 343         return 1;
 344     return 0;
 345 }
 346
 347 #define more_rbsp_data(var) ((var) = cbs_h2645_read_more_rbsp_data(rw))
 348
 349 #define byte_alignment(rw) (get_bits_count(rw) % 8)
 350
 351 #define allocate(name, size) do { \
 352         name ## _ref = av_buffer_allocz(size + \
 353                                         AV_INPUT_BUFFER_PADDING_SIZE); \
 354         if (!name ## _ref) \
 355             return AVERROR(ENOMEM); \
 356         name = name ## _ref->data; \
 357     } while (0)
 358
 359 #define FUNC(name) FUNC_H264(READWRITE, name)
 360 #include "cbs_h264_syntax_template.c"
 361 #undef FUNC
 362
 363 #define FUNC(name) FUNC_H265(READWRITE, name)
 364 #include "cbs_h265_syntax_template.c"
 365 #undef FUNC
 366
 367 #undef READ
 368 #undef READWRITE
 369 #undef RWContext
 370 #undef xu
 371 #undef xi
 372 #undef xue
 373 #undef xse
 374 #undef infer
 375 #undef more_rbsp_data
 376 #undef byte_alignment
 377 #undef allocate
 378
 379
 380 #define WRITE
 381 #define READWRITE write
 382 #define RWContext PutBitContext
 383
 384 #define xu(width, name, var, range_min, range_max, subs, ...) do { \
 385         uint32_t value = var; \
 386         CHECK(ff_cbs_write_unsigned(ctx, rw, width, #name, \
 387                                     SUBSCRIPTS(subs, __VA_ARGS__), \
 388                                     value, range_min, range_max)); \
 389     } while (0)
 390 #define xue(name, var, range_min, range_max, subs, ...) do { \
 391         uint32_t value = var; \
 392         CHECK(cbs_write_ue_golomb(ctx, rw, #name, \
 393                                   SUBSCRIPTS(subs, __VA_ARGS__), \
 394                                   value, range_min, range_max)); \
 395     } while (0)
 396 #define xi(width, name, var, range_min, range_max, subs, ...) do { \
 397         int32_t value = var; \
 398         CHECK(ff_cbs_write_signed(ctx, rw, width, #name, \
 399                                   SUBSCRIPTS(subs, __VA_ARGS__), \
 400                                   value, range_min, range_max)); \
 401     } while (0)
 402 #define xse(name, var, range_min, range_max, subs, ...) do { \
 403         int32_t value = var; \
 404         CHECK(cbs_write_se_golomb(ctx, rw, #name, \
 405                                   SUBSCRIPTS(subs, __VA_ARGS__), \
 406                                   value, range_min, range_max)); \
 407     } while (0)
 408
 409 #define infer(name, value) do { \
 410         if (current->name != (value)) { \
 411             av_log(ctx->log_ctx, AV_LOG_ERROR, \
 412                    "%s does not match inferred value: " \
 413                    "%"PRId64", but should be %"PRId64".\n", \
 414                    #name, (int64_t)current->name, (int64_t)(value)); \
 415             return AVERROR_INVALIDDATA; \
 416         } \
 417     } while (0)
 418
 419 #define more_rbsp_data(var) (var)
 420
 421 #define byte_alignment(rw) (put_bits_count(rw) % 8)
 422
 423 #define allocate(name, size) do { \
 424         if (!name) { \
 425             av_log(ctx->log_ctx, AV_LOG_ERROR, "%s must be set " \
 426                    "for writing.\n", #name); \
 427             return AVERROR_INVALIDDATA; \
 428         } \
 429     } while (0)
 430
 431 #define FUNC(name) FUNC_H264(READWRITE, name)
 432 #include "cbs_h264_syntax_template.c"
 433 #undef FUNC
 434
 435 #define FUNC(name) FUNC_H265(READWRITE, name)
 436 #include "cbs_h265_syntax_template.c"
 437 #undef FUNC
 438
 439 #undef WRITE
 440 #undef READWRITE
 441 #undef RWContext
 442 #undef xu
 443 #undef xi
 444 #undef xue
 445 #undef xse
 446 #undef u
 447 #undef i
 448 #undef flag
 449 #undef ue
 450 #undef se
 451 #undef infer
 452 #undef more_rbsp_data
 453 #undef byte_alignment
 454 #undef allocate
 455
 456
 457 static void cbs_h265_free_vps(void *opaque, uint8_t *content)
 458 {
 459     H265RawVPS *vps = (H265RawVPS*)content;
 460     av_buffer_unref(&vps->extension_data.data_ref);
 461     av_freep(&content);
 462 }
 463
 464 static void cbs_h265_free_sps(void *opaque, uint8_t *content)
 465 {
 466     H265RawSPS *sps = (H265RawSPS*)content;
 467     av_buffer_unref(&sps->extension_data.data_ref);
 468     av_freep(&content);
 469 }
 470
 471 static void cbs_h265_free_pps(void *opaque, uint8_t *content)
 472 {
 473     H265RawPPS *pps = (H265RawPPS*)content;
 474     av_buffer_unref(&pps->extension_data.data_ref);
 475     av_freep(&content);
 476 }
 477
 478 static void cbs_h265_free_slice(void *opaque, uint8_t *content)
 479 {
 480     H265RawSlice *slice = (H265RawSlice*)content;
 481     av_buffer_unref(&slice->data_ref);
 482     av_freep(&content);
 483 }
 484
 485 static void cbs_h265_free_sei_payload(H265RawSEIPayload *payload)
 486 {
 487     switch (payload->payload_type) {
 488     case HEVC_SEI_TYPE_BUFFERING_PERIOD:
 489     case HEVC_SEI_TYPE_PICTURE_TIMING:
 490     case HEVC_SEI_TYPE_PAN_SCAN_RECT:
 491     case HEVC_SEI_TYPE_RECOVERY_POINT:
 492     case HEVC_SEI_TYPE_DISPLAY_ORIENTATION:
 493     case HEVC_SEI_TYPE_ACTIVE_PARAMETER_SETS:
 494     case HEVC_SEI_TYPE_DECODED_PICTURE_HASH:
 495     case HEVC_SEI_TYPE_TIME_CODE:
 496     case HEVC_SEI_TYPE_MASTERING_DISPLAY_INFO:
 497     case HEVC_SEI_TYPE_CONTENT_LIGHT_LEVEL_INFO:
 498     case HEVC_SEI_TYPE_ALTERNATIVE_TRANSFER_CHARACTERISTICS:
 499     case HEVC_SEI_TYPE_ALPHA_CHANNEL_INFO:
 500         break;
 501     case HEVC_SEI_TYPE_USER_DATA_REGISTERED_ITU_T_T35:
 502         av_buffer_unref(&payload->payload.user_data_registered.data_ref);
 503         break;
 504     case HEVC_SEI_TYPE_USER_DATA_UNREGISTERED:
 505         av_buffer_unref(&payload->payload.user_data_unregistered.data_ref);
 506         break;
 507     default:
 508         av_buffer_unref(&payload->payload.other.data_ref);
 509         break;
 510     }
 511     av_buffer_unref(&payload->extension_data.data_ref);
 512 }
 513
 514 static void cbs_h265_free_sei(void *opaque, uint8_t *content)
 515 {
 516     H265RawSEI *sei = (H265RawSEI*)content;
 517     int i;
 518     for (i = 0; i < sei->payload_count; i++)
 519         cbs_h265_free_sei_payload(&sei->payload[i]);
 520     av_freep(&content);
 521 }
 522
 523 static int cbs_h2645_fragment_add_nals(CodedBitstreamContext *ctx,
 524                                        CodedBitstreamFragment *frag,
 525                                        const H2645Packet *packet)
 526 {
 527     int err, i;
 528
 529     for (i = 0; i < packet->nb_nals; i++) {
 530         const H2645NAL *nal = &packet->nals[i];
 531         AVBufferRef *ref;
 532         size_t size = nal->size;
 533
 534         if (nal->nuh_layer_id > 0)
 535             continue;
 536
 537         // Remove trailing zeroes.
 538         while (size > 0 && nal->data[size - 1] == 0)
 539             --size;
 540         if (size == 0) {
 541             av_log(ctx->log_ctx, AV_LOG_VERBOSE, "Discarding empty 0 NAL unit\n");
 542             continue;
 543         }
 544
 545         ref = (nal->data == nal->raw_data) ? frag->data_ref
 546                                            : packet->rbsp.rbsp_buffer_ref;
 547
 548         err = ff_cbs_insert_unit_data(frag, -1, nal->type,
 549                             (uint8_t*)nal->data, size, ref);
 550         if (err < 0)
 551             return err;
 552     }
 553
 554     return 0;
 555 }
 556
 557 static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
 558                                     CodedBitstreamFragment *frag,
 559                                     int header)
 560 {
 561     enum AVCodecID codec_id = ctx->codec->codec_id;
 562     CodedBitstreamH2645Context *priv = ctx->priv_data;
 563     GetByteContext gbc;
 564     int err;
 565
 566     av_assert0(frag->data && frag->nb_units == 0);
 567     if (frag->data_size == 0)
 568         return 0;
 569
 570     if (header && frag->data[0] && codec_id == AV_CODEC_ID_H264) {
 571         // AVCC header.
 572         size_t size, start, end;
 573         int i, count, version;
 574
 575         priv->mp4 = 1;
 576
 577         bytestream2_init(&gbc, frag->data, frag->data_size);
 578
 579         if (bytestream2_get_bytes_left(&gbc) < 6)
 580             return AVERROR_INVALIDDATA;
 581
 582         version = bytestream2_get_byte(&gbc);
 583         if (version != 1) {
 584             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid AVCC header: "
 585                    "first byte %u.\n", version);
 586             return AVERROR_INVALIDDATA;
 587         }
 588
 589         bytestream2_skip(&gbc, 3);
 590         priv->nal_length_size = (bytestream2_get_byte(&gbc) & 3) + 1;
 591
 592         // SPS array.
 593         count = bytestream2_get_byte(&gbc) & 0x1f;
 594         start = bytestream2_tell(&gbc);
 595         for (i = 0; i < count; i++) {
 596             if (bytestream2_get_bytes_left(&gbc) < 2 * (count - i))
 597                 return AVERROR_INVALIDDATA;
 598             size = bytestream2_get_be16(&gbc);
 599             if (bytestream2_get_bytes_left(&gbc) < size)
 600                 return AVERROR_INVALIDDATA;
 601             bytestream2_skip(&gbc, size);
 602         }
 603         end = bytestream2_tell(&gbc);
 604
 605         err = ff_h2645_packet_split(&priv->read_packet,
 606                                     frag->data + start, end - start,
 607                                     ctx->log_ctx, 1, 2, AV_CODEC_ID_H264, 1, 1);
 608         if (err < 0) {
 609             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC SPS array.\n");
 610             return err;
 611         }
 612         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 613         if (err < 0)
 614             return err;
 615
 616         // PPS array.
 617         count = bytestream2_get_byte(&gbc);
 618         start = bytestream2_tell(&gbc);
 619         for (i = 0; i < count; i++) {
 620             if (bytestream2_get_bytes_left(&gbc) < 2 * (count - i))
 621                 return AVERROR_INVALIDDATA;
 622             size = bytestream2_get_be16(&gbc);
 623             if (bytestream2_get_bytes_left(&gbc) < size)
 624                 return AVERROR_INVALIDDATA;
 625             bytestream2_skip(&gbc, size);
 626         }
 627         end = bytestream2_tell(&gbc);
 628
 629         err = ff_h2645_packet_split(&priv->read_packet,
 630                                     frag->data + start, end - start,
 631                                     ctx->log_ctx, 1, 2, AV_CODEC_ID_H264, 1, 1);
 632         if (err < 0) {
 633             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC PPS array.\n");
 634             return err;
 635         }
 636         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 637         if (err < 0)
 638             return err;
 639
 640         if (bytestream2_get_bytes_left(&gbc) > 0) {
 641             av_log(ctx->log_ctx, AV_LOG_WARNING, "%u bytes left at end of AVCC "
 642                    "header.\n", bytestream2_get_bytes_left(&gbc));
 643         }
 644
 645     } else if (header && frag->data[0] && codec_id == AV_CODEC_ID_HEVC) {
 646         // HVCC header.
 647         size_t size, start, end;
 648         int i, j, nb_arrays, nal_unit_type, nb_nals, version;
 649
 650         priv->mp4 = 1;
 651
 652         bytestream2_init(&gbc, frag->data, frag->data_size);
 653
 654         if (bytestream2_get_bytes_left(&gbc) < 23)
 655             return AVERROR_INVALIDDATA;
 656
 657         version = bytestream2_get_byte(&gbc);
 658         if (version != 1) {
 659             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid HVCC header: "
 660                    "first byte %u.\n", version);
 661             return AVERROR_INVALIDDATA;
 662         }
 663
 664         bytestream2_skip(&gbc, 20);
 665         priv->nal_length_size = (bytestream2_get_byte(&gbc) & 3) + 1;
 666
 667         nb_arrays = bytestream2_get_byte(&gbc);
 668         for (i = 0; i < nb_arrays; i++) {
 669             nal_unit_type = bytestream2_get_byte(&gbc) & 0x3f;
 670             nb_nals = bytestream2_get_be16(&gbc);
 671
 672             start = bytestream2_tell(&gbc);
 673             for (j = 0; j < nb_nals; j++) {
 674                 if (bytestream2_get_bytes_left(&gbc) < 2)
 675                     return AVERROR_INVALIDDATA;
 676                 size = bytestream2_get_be16(&gbc);
 677                 if (bytestream2_get_bytes_left(&gbc) < size)
 678                     return AVERROR_INVALIDDATA;
 679                 bytestream2_skip(&gbc, size);
 680             }
 681             end = bytestream2_tell(&gbc);
 682
 683             err = ff_h2645_packet_split(&priv->read_packet,
 684                                         frag->data + start, end - start,
 685                                         ctx->log_ctx, 1, 2, AV_CODEC_ID_HEVC, 1, 1);
 686             if (err < 0) {
 687                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split "
 688                        "HVCC array %d (%d NAL units of type %d).\n",
 689                        i, nb_nals, nal_unit_type);
 690                 return err;
 691             }
 692             err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 693             if (err < 0)
 694                 return err;
 695         }
 696
 697     } else {
 698         // Annex B, or later MP4 with already-known parameters.
 699
 700         err = ff_h2645_packet_split(&priv->read_packet,
 701                                     frag->data, frag->data_size,
 702                                     ctx->log_ctx,
 703                                     priv->mp4, priv->nal_length_size,
 704                                     codec_id, 1, 1);
 705         if (err < 0)
 706             return err;
 707
 708         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 709         if (err < 0)
 710             return err;
 711     }
 712
 713     return 0;
 714 }
 715
 716 #define cbs_h2645_replace_ps(h26n, ps_name, ps_var, id_element) \
 717 static int cbs_h26 ## h26n ## _replace_ ## ps_var(CodedBitstreamContext *ctx, \
 718                                                   CodedBitstreamUnit *unit)  \
 719 { \
 720     CodedBitstreamH26 ## h26n ## Context *priv = ctx->priv_data; \
 721     H26 ## h26n ## Raw ## ps_name *ps_var = unit->content; \
 722     unsigned int id = ps_var->id_element; \
 723     if (id >= FF_ARRAY_ELEMS(priv->ps_var)) { \
 724         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid " #ps_name \
 725                " id : %d.\n", id); \
 726         return AVERROR_INVALIDDATA; \
 727     } \
 728     if (priv->ps_var[id] == priv->active_ ## ps_var) \
 729         priv->active_ ## ps_var = NULL ; \
 730     av_buffer_unref(&priv->ps_var ## _ref[id]); \
 731     if (unit->content_ref) \
 732         priv->ps_var ## _ref[id] = av_buffer_ref(unit->content_ref); \
 733     else \
 734         priv->ps_var ## _ref[id] = av_buffer_alloc(sizeof(*ps_var)); \
 735     if (!priv->ps_var ## _ref[id]) \
 736         return AVERROR(ENOMEM); \
 737     priv->ps_var[id] = (H26 ## h26n ## Raw ## ps_name *)priv->ps_var ## _ref[id]->data; \
 738     if (!unit->content_ref) \
 739         memcpy(priv->ps_var[id], ps_var, sizeof(*ps_var)); \
 740     return 0; \
 741 }
 742
 743 cbs_h2645_replace_ps(4, SPS, sps, seq_parameter_set_id)
 744 cbs_h2645_replace_ps(4, PPS, pps, pic_parameter_set_id)
 745 cbs_h2645_replace_ps(5, VPS, vps, vps_video_parameter_set_id)
 746 cbs_h2645_replace_ps(5, SPS, sps, sps_seq_parameter_set_id)
 747 cbs_h2645_replace_ps(5, PPS, pps, pps_pic_parameter_set_id)
 748
 749 static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
 750                                   CodedBitstreamUnit *unit)
 751 {
 752     GetBitContext gbc;
 753     int err;
 754
 755     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
 756     if (err < 0)
 757         return err;
 758
 759     err = ff_cbs_alloc_unit_content2(ctx, unit);
 760     if (err < 0)
 761         return err;
 762
 763     switch (unit->type) {
 764     case H264_NAL_SPS:
 765         {
 766             H264RawSPS *sps = unit->content;
 767
 768             err = cbs_h264_read_sps(ctx, &gbc, sps);
 769             if (err < 0)
 770                 return err;
 771
 772             err = cbs_h264_replace_sps(ctx, unit);
 773             if (err < 0)
 774                 return err;
 775         }
 776         break;
 777
 778     case H264_NAL_SPS_EXT:
 779         {
 780             err = cbs_h264_read_sps_extension(ctx, &gbc, unit->content);
 781             if (err < 0)
 782                 return err;
 783         }
 784         break;
 785
 786     case H264_NAL_PPS:
 787         {
 788             H264RawPPS *pps = unit->content;
 789
 790             err = cbs_h264_read_pps(ctx, &gbc, pps);
 791             if (err < 0)
 792                 return err;
 793
 794             err = cbs_h264_replace_pps(ctx, unit);
 795             if (err < 0)
 796                 return err;
 797         }
 798         break;
 799
 800     case H264_NAL_SLICE:
 801     case H264_NAL_IDR_SLICE:
 802     case H264_NAL_AUXILIARY_SLICE:
 803         {
 804             H264RawSlice *slice = unit->content;
 805             int pos, len;
 806
 807             err = cbs_h264_read_slice_header(ctx, &gbc, &slice->header);
 808             if (err < 0)
 809                 return err;
 810
 811             if (!cbs_h2645_read_more_rbsp_data(&gbc))
 812                 return AVERROR_INVALIDDATA;
 813
 814             pos = get_bits_count(&gbc);
 815             len = unit->data_size;
 816
 817             slice->data_size = len - pos / 8;
 818             slice->data_ref  = av_buffer_ref(unit->data_ref);
 819             if (!slice->data_ref)
 820                 return AVERROR(ENOMEM);
 821             slice->data = unit->data + pos / 8;
 822             slice->data_bit_start = pos % 8;
 823         }
 824         break;
 825
 826     case H264_NAL_AUD:
 827         {
 828             err = cbs_h264_read_aud(ctx, &gbc, unit->content);
 829             if (err < 0)
 830                 return err;
 831         }
 832         break;
 833
 834     case H264_NAL_SEI:
 835         {
 836             err = cbs_h264_read_sei(ctx, &gbc, unit->content);
 837             if (err < 0)
 838                 return err;
 839         }
 840         break;
 841
 842     case H264_NAL_FILLER_DATA:
 843         {
 844             err = cbs_h264_read_filler(ctx, &gbc, unit->content);
 845             if (err < 0)
 846                 return err;
 847         }
 848         break;
 849
 850     case H264_NAL_END_SEQUENCE:
 851     case H264_NAL_END_STREAM:
 852         {
 853             err = (unit->type == H264_NAL_END_SEQUENCE ?
 854                    cbs_h264_read_end_of_sequence :
 855                    cbs_h264_read_end_of_stream)(ctx, &gbc, unit->content);
 856             if (err < 0)
 857                 return err;
 858         }
 859         break;
 860
 861     default:
 862         return AVERROR(ENOSYS);
 863     }
 864
 865     return 0;
 866 }
 867
 868 static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
 869                                   CodedBitstreamUnit *unit)
 870 {
 871     GetBitContext gbc;
 872     int err;
 873
 874     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
 875     if (err < 0)
 876         return err;
 877
 878     switch (unit->type) {
 879     case HEVC_NAL_VPS:
 880         {
 881             H265RawVPS *vps;
 882
 883             err = ff_cbs_alloc_unit_content(unit, sizeof(*vps),
 884                                             &cbs_h265_free_vps);
 885             if (err < 0)
 886                 return err;
 887             vps = unit->content;
 888
 889             err = cbs_h265_read_vps(ctx, &gbc, vps);
 890             if (err < 0)
 891                 return err;
 892
 893             err = cbs_h265_replace_vps(ctx, unit);
 894             if (err < 0)
 895                 return err;
 896         }
 897         break;
 898     case HEVC_NAL_SPS:
 899         {
 900             H265RawSPS *sps;
 901
 902             err = ff_cbs_alloc_unit_content(unit, sizeof(*sps),
 903                                             &cbs_h265_free_sps);
 904             if (err < 0)
 905                 return err;
 906             sps = unit->content;
 907
 908             err = cbs_h265_read_sps(ctx, &gbc, sps);
 909             if (err < 0)
 910                 return err;
 911
 912             err = cbs_h265_replace_sps(ctx, unit);
 913             if (err < 0)
 914                 return err;
 915         }
 916         break;
 917
 918     case HEVC_NAL_PPS:
 919         {
 920             H265RawPPS *pps;
 921
 922             err = ff_cbs_alloc_unit_content(unit, sizeof(*pps),
 923                                             &cbs_h265_free_pps);
 924             if (err < 0)
 925                 return err;
 926             pps = unit->content;
 927
 928             err = cbs_h265_read_pps(ctx, &gbc, pps);
 929             if (err < 0)
 930                 return err;
 931
 932             err = cbs_h265_replace_pps(ctx, unit);
 933             if (err < 0)
 934                 return err;
 935         }
 936         break;
 937
 938     case HEVC_NAL_TRAIL_N:
 939     case HEVC_NAL_TRAIL_R:
 940     case HEVC_NAL_TSA_N:
 941     case HEVC_NAL_TSA_R:
 942     case HEVC_NAL_STSA_N:
 943     case HEVC_NAL_STSA_R:
 944     case HEVC_NAL_RADL_N:
 945     case HEVC_NAL_RADL_R:
 946     case HEVC_NAL_RASL_N:
 947     case HEVC_NAL_RASL_R:
 948     case HEVC_NAL_BLA_W_LP:
 949     case HEVC_NAL_BLA_W_RADL:
 950     case HEVC_NAL_BLA_N_LP:
 951     case HEVC_NAL_IDR_W_RADL:
 952     case HEVC_NAL_IDR_N_LP:
 953     case HEVC_NAL_CRA_NUT:
 954         {
 955             H265RawSlice *slice;
 956             int pos, len;
 957
 958             err = ff_cbs_alloc_unit_content(unit, sizeof(*slice),
 959                                             &cbs_h265_free_slice);
 960             if (err < 0)
 961                 return err;
 962             slice = unit->content;
 963
 964             err = cbs_h265_read_slice_segment_header(ctx, &gbc, &slice->header);
 965             if (err < 0)
 966                 return err;
 967
 968             if (!cbs_h2645_read_more_rbsp_data(&gbc))
 969                 return AVERROR_INVALIDDATA;
 970
 971             pos = get_bits_count(&gbc);
 972             len = unit->data_size;
 973
 974             slice->data_size = len - pos / 8;
 975             slice->data_ref  = av_buffer_ref(unit->data_ref);
 976             if (!slice->data_ref)
 977                 return AVERROR(ENOMEM);
 978             slice->data = unit->data + pos / 8;
 979             slice->data_bit_start = pos % 8;
 980         }
 981         break;
 982
 983     case HEVC_NAL_AUD:
 984         {
 985             err = ff_cbs_alloc_unit_content(unit,
 986                                             sizeof(H265RawAUD), NULL);
 987             if (err < 0)
 988                 return err;
 989
 990             err = cbs_h265_read_aud(ctx, &gbc, unit->content);
 991             if (err < 0)
 992                 return err;
 993         }
 994         break;
 995
 996     case HEVC_NAL_SEI_PREFIX:
 997     case HEVC_NAL_SEI_SUFFIX:
 998         {
 999             err = ff_cbs_alloc_unit_content(unit, sizeof(H265RawSEI),
1000                                             &cbs_h265_free_sei);
1001
1002             if (err < 0)
1003                 return err;
1004
1005             err = cbs_h265_read_sei(ctx, &gbc, unit->content,
1006                                     unit->type == HEVC_NAL_SEI_PREFIX);
1007
1008             if (err < 0)
1009                 return err;
1010         }
1011         break;
1012
1013     default:
1014         return AVERROR(ENOSYS);
1015     }
1016
1017     return 0;
1018 }
1019
1020 static int cbs_h2645_write_slice_data(CodedBitstreamContext *ctx,
1021                                       PutBitContext *pbc, const uint8_t *data,
1022                                       size_t data_size, int data_bit_start)
1023 {
1024     size_t rest  = data_size - (data_bit_start + 7) / 8;
1025     const uint8_t *pos = data + data_bit_start / 8;
1026
1027     av_assert0(data_bit_start >= 0 &&
1028                data_size > data_bit_start / 8);
1029
1030     if (data_size * 8 + 8 > put_bits_left(pbc))
1031         return AVERROR(ENOSPC);
1032
1033     if (!rest)
1034         goto rbsp_stop_one_bit;
1035
1036     // First copy the remaining bits of the first byte
1037     // The above check ensures that we do not accidentally
1038     // copy beyond the rbsp_stop_one_bit.
1039     if (data_bit_start % 8)
1040         put_bits(pbc, 8 - data_bit_start % 8,
1041                  *pos++ & MAX_UINT_BITS(8 - data_bit_start % 8));
1042
1043     if (put_bits_count(pbc) % 8 == 0) {
1044         // If the writer is aligned at this point,
1045         // memcpy can be used to improve performance.
1046         // This happens normally for CABAC.
1047         flush_put_bits(pbc);
1048         memcpy(put_bits_ptr(pbc), pos, rest);
1049         skip_put_bytes(pbc, rest);
1050     } else {
1051         // If not, we have to copy manually.
1052         // rbsp_stop_one_bit forces us to special-case
1053         // the last byte.
1054         uint8_t temp;
1055         int i;
1056
1057         for (; rest > 4; rest -= 4, pos += 4)
1058             put_bits32(pbc, AV_RB32(pos));
1059
1060         for (; rest > 1; rest--, pos++)
1061             put_bits(pbc, 8, *pos);
1062
1063     rbsp_stop_one_bit:
1064         temp = rest ? *pos : *pos & MAX_UINT_BITS(8 - data_bit_start % 8);
1065
1066         av_assert0(temp);
1067         i = ff_ctz(*pos);
1068         temp = temp >> i;
1069         i = rest ? (8 - i) : (8 - i - data_bit_start % 8);
1070         put_bits(pbc, i, temp);
1071         if (put_bits_count(pbc) % 8)
1072             put_bits(pbc, 8 - put_bits_count(pbc) % 8, 0);
1073     }
1074
1075     return 0;
1076 }
1077
1078 static int cbs_h264_write_nal_unit(CodedBitstreamContext *ctx,
1079                                    CodedBitstreamUnit *unit,
1080                                    PutBitContext *pbc)
1081 {
1082     int err;
1083
1084     switch (unit->type) {
1085     case H264_NAL_SPS:
1086         {
1087             H264RawSPS *sps = unit->content;
1088
1089             err = cbs_h264_write_sps(ctx, pbc, sps);
1090             if (err < 0)
1091                 return err;
1092
1093             err = cbs_h264_replace_sps(ctx, unit);
1094             if (err < 0)
1095                 return err;
1096         }
1097         break;
1098
1099     case H264_NAL_SPS_EXT:
1100         {
1101             H264RawSPSExtension *sps_ext = unit->content;
1102
1103             err = cbs_h264_write_sps_extension(ctx, pbc, sps_ext);
1104             if (err < 0)
1105                 return err;
1106         }
1107         break;
1108
1109     case H264_NAL_PPS:
1110         {
1111             H264RawPPS *pps = unit->content;
1112
1113             err = cbs_h264_write_pps(ctx, pbc, pps);
1114             if (err < 0)
1115                 return err;
1116
1117             err = cbs_h264_replace_pps(ctx, unit);
1118             if (err < 0)
1119                 return err;
1120         }
1121         break;
1122
1123     case H264_NAL_SLICE:
1124     case H264_NAL_IDR_SLICE:
1125     case H264_NAL_AUXILIARY_SLICE:
1126         {
1127             H264RawSlice *slice = unit->content;
1128
1129             err = cbs_h264_write_slice_header(ctx, pbc, &slice->header);
1130             if (err < 0)
1131                 return err;
1132
1133             if (slice->data) {
1134                 err = cbs_h2645_write_slice_data(ctx, pbc, slice->data,
1135                                                  slice->data_size,
1136                                                  slice->data_bit_start);
1137                 if (err < 0)
1138                     return err;
1139             } else {
1140                 // No slice data - that was just the header.
1141                 // (Bitstream may be unaligned!)
1142             }
1143         }
1144         break;
1145
1146     case H264_NAL_AUD:
1147         {
1148             err = cbs_h264_write_aud(ctx, pbc, unit->content);
1149             if (err < 0)
1150                 return err;
1151         }
1152         break;
1153
1154     case H264_NAL_SEI:
1155         {
1156             err = cbs_h264_write_sei(ctx, pbc, unit->content);
1157             if (err < 0)
1158                 return err;
1159         }
1160         break;
1161
1162     case H264_NAL_FILLER_DATA:
1163         {
1164             err = cbs_h264_write_filler(ctx, pbc, unit->content);
1165             if (err < 0)
1166                 return err;
1167         }
1168         break;
1169
1170     case H264_NAL_END_SEQUENCE:
1171         {
1172             err = cbs_h264_write_end_of_sequence(ctx, pbc, unit->content);
1173             if (err < 0)
1174                 return err;
1175         }
1176         break;
1177
1178     case H264_NAL_END_STREAM:
1179         {
1180             err = cbs_h264_write_end_of_stream(ctx, pbc, unit->content);
1181             if (err < 0)
1182                 return err;
1183         }
1184         break;
1185
1186     default:
1187         av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for "
1188                "NAL unit type %"PRIu32".\n", unit->type);
1189         return AVERROR_PATCHWELCOME;
1190     }
1191
1192     return 0;
1193 }
1194
1195 static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
1196                                    CodedBitstreamUnit *unit,
1197                                    PutBitContext *pbc)
1198 {
1199     int err;
1200
1201     switch (unit->type) {
1202     case HEVC_NAL_VPS:
1203         {
1204             H265RawVPS *vps = unit->content;
1205
1206             err = cbs_h265_write_vps(ctx, pbc, vps);
1207             if (err < 0)
1208                 return err;
1209
1210             err = cbs_h265_replace_vps(ctx, unit);
1211             if (err < 0)
1212                 return err;
1213         }
1214         break;
1215
1216     case HEVC_NAL_SPS:
1217         {
1218             H265RawSPS *sps = unit->content;
1219
1220             err = cbs_h265_write_sps(ctx, pbc, sps);
1221             if (err < 0)
1222                 return err;
1223
1224             err = cbs_h265_replace_sps(ctx, unit);
1225             if (err < 0)
1226                 return err;
1227         }
1228         break;
1229
1230     case HEVC_NAL_PPS:
1231         {
1232             H265RawPPS *pps = unit->content;
1233
1234             err = cbs_h265_write_pps(ctx, pbc, pps);
1235             if (err < 0)
1236                 return err;
1237
1238             err = cbs_h265_replace_pps(ctx, unit);
1239             if (err < 0)
1240                 return err;
1241         }
1242         break;
1243
1244     case HEVC_NAL_TRAIL_N:
1245     case HEVC_NAL_TRAIL_R:
1246     case HEVC_NAL_TSA_N:
1247     case HEVC_NAL_TSA_R:
1248     case HEVC_NAL_STSA_N:
1249     case HEVC_NAL_STSA_R:
1250     case HEVC_NAL_RADL_N:
1251     case HEVC_NAL_RADL_R:
1252     case HEVC_NAL_RASL_N:
1253     case HEVC_NAL_RASL_R:
1254     case HEVC_NAL_BLA_W_LP:
1255     case HEVC_NAL_BLA_W_RADL:
1256     case HEVC_NAL_BLA_N_LP:
1257     case HEVC_NAL_IDR_W_RADL:
1258     case HEVC_NAL_IDR_N_LP:
1259     case HEVC_NAL_CRA_NUT:
1260         {
1261             H265RawSlice *slice = unit->content;
1262
1263             err = cbs_h265_write_slice_segment_header(ctx, pbc, &slice->header);
1264             if (err < 0)
1265                 return err;
1266
1267             if (slice->data) {
1268                 err = cbs_h2645_write_slice_data(ctx, pbc, slice->data,
1269                                                  slice->data_size,
1270                                                  slice->data_bit_start);
1271                 if (err < 0)
1272                     return err;
1273             } else {
1274                 // No slice data - that was just the header.
1275             }
1276         }
1277         break;
1278
1279     case HEVC_NAL_AUD:
1280         {
1281             err = cbs_h265_write_aud(ctx, pbc, unit->content);
1282             if (err < 0)
1283                 return err;
1284         }
1285         break;
1286
1287     case HEVC_NAL_SEI_PREFIX:
1288     case HEVC_NAL_SEI_SUFFIX:
1289         {
1290             err = cbs_h265_write_sei(ctx, pbc, unit->content,
1291                                      unit->type == HEVC_NAL_SEI_PREFIX);
1292
1293             if (err < 0)
1294                 return err;
1295         }
1296         break;
1297
1298     default:
1299         av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for "
1300                "NAL unit type %"PRIu32".\n", unit->type);
1301         return AVERROR_PATCHWELCOME;
1302     }
1303
1304     return 0;
1305 }
1306
1307 static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
1308                                        CodedBitstreamFragment *frag)
1309 {
1310     uint8_t *data;
1311     size_t max_size, dp, sp;
1312     int err, i, zero_run;
1313
1314     for (i = 0; i < frag->nb_units; i++) {
1315         // Data should already all have been written when we get here.
1316         av_assert0(frag->units[i].data);
1317     }
1318
1319     max_size = 0;
1320     for (i = 0; i < frag->nb_units; i++) {
1321         // Start code + content with worst-case emulation prevention.
1322         max_size += 4 + frag->units[i].data_size * 3 / 2;
1323     }
1324
1325     data = av_realloc(NULL, max_size + AV_INPUT_BUFFER_PADDING_SIZE);
1326     if (!data)
1327         return AVERROR(ENOMEM);
1328
1329     dp = 0;
1330     for (i = 0; i < frag->nb_units; i++) {
1331         CodedBitstreamUnit *unit = &frag->units[i];
1332
1333         if (unit->data_bit_padding > 0) {
1334             if (i < frag->nb_units - 1)
1335                 av_log(ctx->log_ctx, AV_LOG_WARNING, "Probably invalid "
1336                        "unaligned padding on non-final NAL unit.\n");
1337             else
1338                 frag->data_bit_padding = unit->data_bit_padding;
1339         }
1340
1341         if ((ctx->codec->codec_id == AV_CODEC_ID_H264 &&
1342              (unit->type == H264_NAL_SPS ||
1343               unit->type == H264_NAL_PPS)) ||
1344             (ctx->codec->codec_id == AV_CODEC_ID_HEVC &&
1345              (unit->type == HEVC_NAL_VPS ||
1346               unit->type == HEVC_NAL_SPS ||
1347               unit->type == HEVC_NAL_PPS)) ||
1348             i == 0 /* (Assume this is the start of an access unit.) */) {
1349             // zero_byte
1350             data[dp++] = 0;
1351         }
1352         // start_code_prefix_one_3bytes
1353         data[dp++] = 0;
1354         data[dp++] = 0;
1355         data[dp++] = 1;
1356
1357         zero_run = 0;
1358         for (sp = 0; sp < unit->data_size; sp++) {
1359             if (zero_run < 2) {
1360                 if (unit->data[sp] == 0)
1361                     ++zero_run;
1362                 else
1363                     zero_run = 0;
1364             } else {
1365                 if ((unit->data[sp] & ~3) == 0) {
1366                     // emulation_prevention_three_byte
1367                     data[dp++] = 3;
1368                 }
1369                 zero_run = unit->data[sp] == 0;
1370             }
1371             data[dp++] = unit->data[sp];
1372         }
1373     }
1374
1375     av_assert0(dp <= max_size);
1376     err = av_reallocp(&data, dp + AV_INPUT_BUFFER_PADDING_SIZE);
1377     if (err)
1378         return err;
1379     memset(data + dp, 0, AV_INPUT_BUFFER_PADDING_SIZE);
1380
1381     frag->data_ref = av_buffer_create(data, dp + AV_INPUT_BUFFER_PADDING_SIZE,
1382                                       NULL, NULL, 0);
1383     if (!frag->data_ref) {
1384         av_freep(&data);
1385         return AVERROR(ENOMEM);
1386     }
1387
1388     frag->data = data;
1389     frag->data_size = dp;
1390
1391     return 0;
1392 }
1393
1394 static void cbs_h264_close(CodedBitstreamContext *ctx)
1395 {
1396     CodedBitstreamH264Context *h264 = ctx->priv_data;
1397     int i;
1398
1399     ff_h2645_packet_uninit(&h264->common.read_packet);
1400
1401     for (i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++)
1402         av_buffer_unref(&h264->sps_ref[i]);
1403     for (i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++)
1404         av_buffer_unref(&h264->pps_ref[i]);
1405 }
1406
1407 static void cbs_h265_close(CodedBitstreamContext *ctx)
1408 {
1409     CodedBitstreamH265Context *h265 = ctx->priv_data;
1410     int i;
1411
1412     ff_h2645_packet_uninit(&h265->common.read_packet);
1413
1414     for (i = 0; i < FF_ARRAY_ELEMS(h265->vps); i++)
1415         av_buffer_unref(&h265->vps_ref[i]);
1416     for (i = 0; i < FF_ARRAY_ELEMS(h265->sps); i++)
1417         av_buffer_unref(&h265->sps_ref[i]);
1418     for (i = 0; i < FF_ARRAY_ELEMS(h265->pps); i++)
1419         av_buffer_unref(&h265->pps_ref[i]);
1420 }
1421
1422 static void cbs_h264_free_sei_payload(H264RawSEIPayload *payload)
1423 {
1424     switch (payload->payload_type) {
1425     case H264_SEI_TYPE_BUFFERING_PERIOD:
1426     case H264_SEI_TYPE_PIC_TIMING:
1427     case H264_SEI_TYPE_PAN_SCAN_RECT:
1428     case H264_SEI_TYPE_RECOVERY_POINT:
1429     case H264_SEI_TYPE_DISPLAY_ORIENTATION:
1430     case H264_SEI_TYPE_MASTERING_DISPLAY_COLOUR_VOLUME:
1431     case H264_SEI_TYPE_ALTERNATIVE_TRANSFER:
1432         break;
1433     case H264_SEI_TYPE_USER_DATA_REGISTERED:
1434         av_buffer_unref(&payload->payload.user_data_registered.data_ref);
1435         break;
1436     case H264_SEI_TYPE_USER_DATA_UNREGISTERED:
1437         av_buffer_unref(&payload->payload.user_data_unregistered.data_ref);
1438         break;
1439     default:
1440         av_buffer_unref(&payload->payload.other.data_ref);
1441         break;
1442     }
1443 }
1444
1445 static void cbs_h264_free_sei(void *opaque, uint8_t *content)
1446 {
1447     H264RawSEI *sei = (H264RawSEI*)content;
1448     int i;
1449     for (i = 0; i < sei->payload_count; i++)
1450         cbs_h264_free_sei_payload(&sei->payload[i]);
1451     av_freep(&content);
1452 }
1453
1454 static const CodedBitstreamUnitTypeDescriptor cbs_h264_unit_types[] = {
1455     CBS_UNIT_TYPE_POD(H264_NAL_SPS,     H264RawSPS),
1456     CBS_UNIT_TYPE_POD(H264_NAL_SPS_EXT, H264RawSPSExtension),
1457
1458     CBS_UNIT_TYPE_INTERNAL_REF(H264_NAL_PPS, H264RawPPS, slice_group_id),
1459
1460     {
1461         .nb_unit_types  = 3,
1462         .unit_types     = {
1463             H264_NAL_IDR_SLICE,
1464             H264_NAL_SLICE,
1465             H264_NAL_AUXILIARY_SLICE,
1466         },
1467         .content_type   = CBS_CONTENT_TYPE_INTERNAL_REFS,
1468         .content_size   = sizeof(H264RawSlice),
1469         .nb_ref_offsets = 1,
1470         .ref_offsets    = { offsetof(H264RawSlice, data) },
1471     },
1472
1473     CBS_UNIT_TYPE_POD(H264_NAL_AUD,          H264RawAUD),
1474     CBS_UNIT_TYPE_POD(H264_NAL_FILLER_DATA,  H264RawFiller),
1475     CBS_UNIT_TYPE_POD(H264_NAL_END_SEQUENCE, H264RawNALUnitHeader),
1476     CBS_UNIT_TYPE_POD(H264_NAL_END_STREAM,   H264RawNALUnitHeader),
1477
1478     CBS_UNIT_TYPE_COMPLEX(H264_NAL_SEI, H264RawSEI, &cbs_h264_free_sei),
1479
1480     CBS_UNIT_TYPE_END_OF_LIST
1481 };
1482
1483 const CodedBitstreamType ff_cbs_type_h264 = {
1484     .codec_id          = AV_CODEC_ID_H264,
1485
1486     .priv_data_size    = sizeof(CodedBitstreamH264Context),
1487
1488     .unit_types        = cbs_h264_unit_types,
1489
1490     .split_fragment    = &cbs_h2645_split_fragment,
1491     .read_unit         = &cbs_h264_read_nal_unit,
1492     .write_unit        = &cbs_h264_write_nal_unit,
1493     .assemble_fragment = &cbs_h2645_assemble_fragment,
1494
1495     .close             = &cbs_h264_close,
1496 };
1497
1498 const CodedBitstreamType ff_cbs_type_h265 = {
1499     .codec_id          = AV_CODEC_ID_HEVC,
1500
1501     .priv_data_size    = sizeof(CodedBitstreamH265Context),
1502
1503     .split_fragment    = &cbs_h2645_split_fragment,
1504     .read_unit         = &cbs_h265_read_nal_unit,
1505     .write_unit        = &cbs_h265_write_nal_unit,
1506     .assemble_fragment = &cbs_h2645_assemble_fragment,
1507
1508     .close             = &cbs_h265_close,
1509 };
1510
1511 int ff_cbs_h264_add_sei_message(CodedBitstreamFragment *au,
1512                                 H264RawSEIPayload *payload)
1513 {
1514     H264RawSEI *sei = NULL;
1515     int err, i;
1516
1517     // Find an existing SEI NAL unit to add to.
1518     for (i = 0; i < au->nb_units; i++) {
1519         if (au->units[i].type == H264_NAL_SEI) {
1520             sei = au->units[i].content;
1521             if (sei->payload_count < H264_MAX_SEI_PAYLOADS)
1522                 break;
1523
1524             sei = NULL;
1525         }
1526     }
1527
1528     if (!sei) {
1529         // Need to make a new SEI NAL unit.  Insert it before the first
1530         // slice data NAL unit; if no slice data, add at the end.
1531         AVBufferRef *sei_ref;
1532
1533         sei = av_mallocz(sizeof(*sei));
1534         if (!sei) {
1535             err = AVERROR(ENOMEM);
1536             goto fail;
1537         }
1538
1539         sei->nal_unit_header.nal_unit_type = H264_NAL_SEI;
1540         sei->nal_unit_header.nal_ref_idc   = 0;
1541
1542         sei_ref = av_buffer_create((uint8_t*)sei, sizeof(*sei),
1543                                    &cbs_h264_free_sei, NULL, 0);
1544         if (!sei_ref) {
1545             av_freep(&sei);
1546             err = AVERROR(ENOMEM);
1547             goto fail;
1548         }
1549
1550         for (i = 0; i < au->nb_units; i++) {
1551             if (au->units[i].type == H264_NAL_SLICE ||
1552                 au->units[i].type == H264_NAL_IDR_SLICE)
1553                 break;
1554         }
1555
1556         err = ff_cbs_insert_unit_content(au, i, H264_NAL_SEI,
1557                                          sei, sei_ref);
1558         av_buffer_unref(&sei_ref);
1559         if (err < 0)
1560             goto fail;
1561     }
1562
1563     memcpy(&sei->payload[sei->payload_count], payload, sizeof(*payload));
1564     ++sei->payload_count;
1565
1566     return 0;
1567 fail:
1568     cbs_h264_free_sei_payload(payload);
1569     return err;
1570 }
1571
1572 void ff_cbs_h264_delete_sei_message(CodedBitstreamFragment *au,
1573                                     CodedBitstreamUnit *nal,
1574                                     int position)
1575 {
1576     H264RawSEI *sei = nal->content;
1577
1578     av_assert0(nal->type == H264_NAL_SEI);
1579     av_assert0(position >= 0 && position < sei->payload_count);
1580
1581     if (position == 0 && sei->payload_count == 1) {
1582         // Deleting NAL unit entirely.
1583         int i;
1584
1585         for (i = 0; i < au->nb_units; i++) {
1586             if (&au->units[i] == nal)
1587                 break;
1588         }
1589
1590         ff_cbs_delete_unit(au, i);
1591     } else {
1592         cbs_h264_free_sei_payload(&sei->payload[position]);
1593
1594         --sei->payload_count;
1595         memmove(sei->payload + position,
1596                 sei->payload + position + 1,
1597                 (sei->payload_count - position) * sizeof(*sei->payload));
1598     }
1599 }