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 int cbs_h2645_fragment_add_nals(CodedBitstreamContext *ctx,
 458                                        CodedBitstreamFragment *frag,
 459                                        const H2645Packet *packet)
 460 {
 461     int err, i;
 462
 463     for (i = 0; i < packet->nb_nals; i++) {
 464         const H2645NAL *nal = &packet->nals[i];
 465         AVBufferRef *ref;
 466         size_t size = nal->size;
 467
 468         if (nal->nuh_layer_id > 0)
 469             continue;
 470
 471         // Remove trailing zeroes.
 472         while (size > 0 && nal->data[size - 1] == 0)
 473             --size;
 474         if (size == 0) {
 475             av_log(ctx->log_ctx, AV_LOG_VERBOSE, "Discarding empty 0 NAL unit\n");
 476             continue;
 477         }
 478
 479         ref = (nal->data == nal->raw_data) ? frag->data_ref
 480                                            : packet->rbsp.rbsp_buffer_ref;
 481
 482         err = ff_cbs_insert_unit_data(frag, -1, nal->type,
 483                             (uint8_t*)nal->data, size, ref);
 484         if (err < 0)
 485             return err;
 486     }
 487
 488     return 0;
 489 }
 490
 491 static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
 492                                     CodedBitstreamFragment *frag,
 493                                     int header)
 494 {
 495     enum AVCodecID codec_id = ctx->codec->codec_id;
 496     CodedBitstreamH2645Context *priv = ctx->priv_data;
 497     GetByteContext gbc;
 498     int err;
 499
 500     av_assert0(frag->data && frag->nb_units == 0);
 501     if (frag->data_size == 0)
 502         return 0;
 503
 504     if (header && frag->data[0] && codec_id == AV_CODEC_ID_H264) {
 505         // AVCC header.
 506         size_t size, start, end;
 507         int i, count, version;
 508
 509         priv->mp4 = 1;
 510
 511         bytestream2_init(&gbc, frag->data, frag->data_size);
 512
 513         if (bytestream2_get_bytes_left(&gbc) < 6)
 514             return AVERROR_INVALIDDATA;
 515
 516         version = bytestream2_get_byte(&gbc);
 517         if (version != 1) {
 518             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid AVCC header: "
 519                    "first byte %u.\n", version);
 520             return AVERROR_INVALIDDATA;
 521         }
 522
 523         bytestream2_skip(&gbc, 3);
 524         priv->nal_length_size = (bytestream2_get_byte(&gbc) & 3) + 1;
 525
 526         // SPS array.
 527         count = bytestream2_get_byte(&gbc) & 0x1f;
 528         start = bytestream2_tell(&gbc);
 529         for (i = 0; i < count; i++) {
 530             if (bytestream2_get_bytes_left(&gbc) < 2 * (count - i))
 531                 return AVERROR_INVALIDDATA;
 532             size = bytestream2_get_be16(&gbc);
 533             if (bytestream2_get_bytes_left(&gbc) < size)
 534                 return AVERROR_INVALIDDATA;
 535             bytestream2_skip(&gbc, size);
 536         }
 537         end = bytestream2_tell(&gbc);
 538
 539         err = ff_h2645_packet_split(&priv->read_packet,
 540                                     frag->data + start, end - start,
 541                                     ctx->log_ctx, 1, 2, AV_CODEC_ID_H264, 1, 1);
 542         if (err < 0) {
 543             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC SPS array.\n");
 544             return err;
 545         }
 546         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 547         if (err < 0)
 548             return err;
 549
 550         // PPS array.
 551         count = bytestream2_get_byte(&gbc);
 552         start = bytestream2_tell(&gbc);
 553         for (i = 0; i < count; i++) {
 554             if (bytestream2_get_bytes_left(&gbc) < 2 * (count - i))
 555                 return AVERROR_INVALIDDATA;
 556             size = bytestream2_get_be16(&gbc);
 557             if (bytestream2_get_bytes_left(&gbc) < size)
 558                 return AVERROR_INVALIDDATA;
 559             bytestream2_skip(&gbc, size);
 560         }
 561         end = bytestream2_tell(&gbc);
 562
 563         err = ff_h2645_packet_split(&priv->read_packet,
 564                                     frag->data + start, end - start,
 565                                     ctx->log_ctx, 1, 2, AV_CODEC_ID_H264, 1, 1);
 566         if (err < 0) {
 567             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC PPS array.\n");
 568             return err;
 569         }
 570         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 571         if (err < 0)
 572             return err;
 573
 574         if (bytestream2_get_bytes_left(&gbc) > 0) {
 575             av_log(ctx->log_ctx, AV_LOG_WARNING, "%u bytes left at end of AVCC "
 576                    "header.\n", bytestream2_get_bytes_left(&gbc));
 577         }
 578
 579     } else if (header && frag->data[0] && codec_id == AV_CODEC_ID_HEVC) {
 580         // HVCC header.
 581         size_t size, start, end;
 582         int i, j, nb_arrays, nal_unit_type, nb_nals, version;
 583
 584         priv->mp4 = 1;
 585
 586         bytestream2_init(&gbc, frag->data, frag->data_size);
 587
 588         if (bytestream2_get_bytes_left(&gbc) < 23)
 589             return AVERROR_INVALIDDATA;
 590
 591         version = bytestream2_get_byte(&gbc);
 592         if (version != 1) {
 593             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid HVCC header: "
 594                    "first byte %u.\n", version);
 595             return AVERROR_INVALIDDATA;
 596         }
 597
 598         bytestream2_skip(&gbc, 20);
 599         priv->nal_length_size = (bytestream2_get_byte(&gbc) & 3) + 1;
 600
 601         nb_arrays = bytestream2_get_byte(&gbc);
 602         for (i = 0; i < nb_arrays; i++) {
 603             nal_unit_type = bytestream2_get_byte(&gbc) & 0x3f;
 604             nb_nals = bytestream2_get_be16(&gbc);
 605
 606             start = bytestream2_tell(&gbc);
 607             for (j = 0; j < nb_nals; j++) {
 608                 if (bytestream2_get_bytes_left(&gbc) < 2)
 609                     return AVERROR_INVALIDDATA;
 610                 size = bytestream2_get_be16(&gbc);
 611                 if (bytestream2_get_bytes_left(&gbc) < size)
 612                     return AVERROR_INVALIDDATA;
 613                 bytestream2_skip(&gbc, size);
 614             }
 615             end = bytestream2_tell(&gbc);
 616
 617             err = ff_h2645_packet_split(&priv->read_packet,
 618                                         frag->data + start, end - start,
 619                                         ctx->log_ctx, 1, 2, AV_CODEC_ID_HEVC, 1, 1);
 620             if (err < 0) {
 621                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split "
 622                        "HVCC array %d (%d NAL units of type %d).\n",
 623                        i, nb_nals, nal_unit_type);
 624                 return err;
 625             }
 626             err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 627             if (err < 0)
 628                 return err;
 629         }
 630
 631     } else {
 632         // Annex B, or later MP4 with already-known parameters.
 633
 634         err = ff_h2645_packet_split(&priv->read_packet,
 635                                     frag->data, frag->data_size,
 636                                     ctx->log_ctx,
 637                                     priv->mp4, priv->nal_length_size,
 638                                     codec_id, 1, 1);
 639         if (err < 0)
 640             return err;
 641
 642         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 643         if (err < 0)
 644             return err;
 645     }
 646
 647     return 0;
 648 }
 649
 650 #define cbs_h2645_replace_ps(h26n, ps_name, ps_var, id_element) \
 651 static int cbs_h26 ## h26n ## _replace_ ## ps_var(CodedBitstreamContext *ctx, \
 652                                                   CodedBitstreamUnit *unit)  \
 653 { \
 654     CodedBitstreamH26 ## h26n ## Context *priv = ctx->priv_data; \
 655     H26 ## h26n ## Raw ## ps_name *ps_var = unit->content; \
 656     unsigned int id = ps_var->id_element; \
 657     int err; \
 658     if (id >= FF_ARRAY_ELEMS(priv->ps_var)) { \
 659         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid " #ps_name \
 660                " id : %d.\n", id); \
 661         return AVERROR_INVALIDDATA; \
 662     } \
 663     err = ff_cbs_make_unit_refcounted(ctx, unit); \
 664     if (err < 0) \
 665         return err; \
 666     if (priv->ps_var[id] == priv->active_ ## ps_var) \
 667         priv->active_ ## ps_var = NULL ; \
 668     av_buffer_unref(&priv->ps_var ## _ref[id]); \
 669     av_assert0(unit->content_ref); \
 670     priv->ps_var ## _ref[id] = av_buffer_ref(unit->content_ref); \
 671     if (!priv->ps_var ## _ref[id]) \
 672         return AVERROR(ENOMEM); \
 673     priv->ps_var[id] = (H26 ## h26n ## Raw ## ps_name *)priv->ps_var ## _ref[id]->data; \
 674     return 0; \
 675 }
 676
 677 cbs_h2645_replace_ps(4, SPS, sps, seq_parameter_set_id)
 678 cbs_h2645_replace_ps(4, PPS, pps, pic_parameter_set_id)
 679 cbs_h2645_replace_ps(5, VPS, vps, vps_video_parameter_set_id)
 680 cbs_h2645_replace_ps(5, SPS, sps, sps_seq_parameter_set_id)
 681 cbs_h2645_replace_ps(5, PPS, pps, pps_pic_parameter_set_id)
 682
 683 static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
 684                                   CodedBitstreamUnit *unit)
 685 {
 686     GetBitContext gbc;
 687     int err;
 688
 689     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
 690     if (err < 0)
 691         return err;
 692
 693     err = ff_cbs_alloc_unit_content2(ctx, unit);
 694     if (err < 0)
 695         return err;
 696
 697     switch (unit->type) {
 698     case H264_NAL_SPS:
 699         {
 700             H264RawSPS *sps = unit->content;
 701
 702             err = cbs_h264_read_sps(ctx, &gbc, sps);
 703             if (err < 0)
 704                 return err;
 705
 706             err = cbs_h264_replace_sps(ctx, unit);
 707             if (err < 0)
 708                 return err;
 709         }
 710         break;
 711
 712     case H264_NAL_SPS_EXT:
 713         {
 714             err = cbs_h264_read_sps_extension(ctx, &gbc, unit->content);
 715             if (err < 0)
 716                 return err;
 717         }
 718         break;
 719
 720     case H264_NAL_PPS:
 721         {
 722             H264RawPPS *pps = unit->content;
 723
 724             err = cbs_h264_read_pps(ctx, &gbc, pps);
 725             if (err < 0)
 726                 return err;
 727
 728             err = cbs_h264_replace_pps(ctx, unit);
 729             if (err < 0)
 730                 return err;
 731         }
 732         break;
 733
 734     case H264_NAL_SLICE:
 735     case H264_NAL_IDR_SLICE:
 736     case H264_NAL_AUXILIARY_SLICE:
 737         {
 738             H264RawSlice *slice = unit->content;
 739             int pos, len;
 740
 741             err = cbs_h264_read_slice_header(ctx, &gbc, &slice->header);
 742             if (err < 0)
 743                 return err;
 744
 745             if (!cbs_h2645_read_more_rbsp_data(&gbc))
 746                 return AVERROR_INVALIDDATA;
 747
 748             pos = get_bits_count(&gbc);
 749             len = unit->data_size;
 750
 751             slice->data_size = len - pos / 8;
 752             slice->data_ref  = av_buffer_ref(unit->data_ref);
 753             if (!slice->data_ref)
 754                 return AVERROR(ENOMEM);
 755             slice->data = unit->data + pos / 8;
 756             slice->data_bit_start = pos % 8;
 757         }
 758         break;
 759
 760     case H264_NAL_AUD:
 761         {
 762             err = cbs_h264_read_aud(ctx, &gbc, unit->content);
 763             if (err < 0)
 764                 return err;
 765         }
 766         break;
 767
 768     case H264_NAL_SEI:
 769         {
 770             err = cbs_h264_read_sei(ctx, &gbc, unit->content);
 771             if (err < 0)
 772                 return err;
 773         }
 774         break;
 775
 776     case H264_NAL_FILLER_DATA:
 777         {
 778             err = cbs_h264_read_filler(ctx, &gbc, unit->content);
 779             if (err < 0)
 780                 return err;
 781         }
 782         break;
 783
 784     case H264_NAL_END_SEQUENCE:
 785     case H264_NAL_END_STREAM:
 786         {
 787             err = (unit->type == H264_NAL_END_SEQUENCE ?
 788                    cbs_h264_read_end_of_sequence :
 789                    cbs_h264_read_end_of_stream)(ctx, &gbc, unit->content);
 790             if (err < 0)
 791                 return err;
 792         }
 793         break;
 794
 795     default:
 796         return AVERROR(ENOSYS);
 797     }
 798
 799     return 0;
 800 }
 801
 802 static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
 803                                   CodedBitstreamUnit *unit)
 804 {
 805     GetBitContext gbc;
 806     int err;
 807
 808     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
 809     if (err < 0)
 810         return err;
 811
 812     err = ff_cbs_alloc_unit_content2(ctx, unit);
 813     if (err < 0)
 814         return err;
 815
 816     switch (unit->type) {
 817     case HEVC_NAL_VPS:
 818         {
 819             H265RawVPS *vps = unit->content;
 820
 821             err = cbs_h265_read_vps(ctx, &gbc, vps);
 822             if (err < 0)
 823                 return err;
 824
 825             err = cbs_h265_replace_vps(ctx, unit);
 826             if (err < 0)
 827                 return err;
 828         }
 829         break;
 830     case HEVC_NAL_SPS:
 831         {
 832             H265RawSPS *sps = unit->content;
 833
 834             err = cbs_h265_read_sps(ctx, &gbc, sps);
 835             if (err < 0)
 836                 return err;
 837
 838             err = cbs_h265_replace_sps(ctx, unit);
 839             if (err < 0)
 840                 return err;
 841         }
 842         break;
 843
 844     case HEVC_NAL_PPS:
 845         {
 846             H265RawPPS *pps = unit->content;
 847
 848             err = cbs_h265_read_pps(ctx, &gbc, pps);
 849             if (err < 0)
 850                 return err;
 851
 852             err = cbs_h265_replace_pps(ctx, unit);
 853             if (err < 0)
 854                 return err;
 855         }
 856         break;
 857
 858     case HEVC_NAL_TRAIL_N:
 859     case HEVC_NAL_TRAIL_R:
 860     case HEVC_NAL_TSA_N:
 861     case HEVC_NAL_TSA_R:
 862     case HEVC_NAL_STSA_N:
 863     case HEVC_NAL_STSA_R:
 864     case HEVC_NAL_RADL_N:
 865     case HEVC_NAL_RADL_R:
 866     case HEVC_NAL_RASL_N:
 867     case HEVC_NAL_RASL_R:
 868     case HEVC_NAL_BLA_W_LP:
 869     case HEVC_NAL_BLA_W_RADL:
 870     case HEVC_NAL_BLA_N_LP:
 871     case HEVC_NAL_IDR_W_RADL:
 872     case HEVC_NAL_IDR_N_LP:
 873     case HEVC_NAL_CRA_NUT:
 874         {
 875             H265RawSlice *slice = unit->content;
 876             int pos, len;
 877
 878             err = cbs_h265_read_slice_segment_header(ctx, &gbc, &slice->header);
 879             if (err < 0)
 880                 return err;
 881
 882             if (!cbs_h2645_read_more_rbsp_data(&gbc))
 883                 return AVERROR_INVALIDDATA;
 884
 885             pos = get_bits_count(&gbc);
 886             len = unit->data_size;
 887
 888             slice->data_size = len - pos / 8;
 889             slice->data_ref  = av_buffer_ref(unit->data_ref);
 890             if (!slice->data_ref)
 891                 return AVERROR(ENOMEM);
 892             slice->data = unit->data + pos / 8;
 893             slice->data_bit_start = pos % 8;
 894         }
 895         break;
 896
 897     case HEVC_NAL_AUD:
 898         {
 899             err = cbs_h265_read_aud(ctx, &gbc, unit->content);
 900             if (err < 0)
 901                 return err;
 902         }
 903         break;
 904
 905     case HEVC_NAL_SEI_PREFIX:
 906     case HEVC_NAL_SEI_SUFFIX:
 907         {
 908             err = cbs_h265_read_sei(ctx, &gbc, unit->content,
 909                                     unit->type == HEVC_NAL_SEI_PREFIX);
 910
 911             if (err < 0)
 912                 return err;
 913         }
 914         break;
 915
 916     default:
 917         return AVERROR(ENOSYS);
 918     }
 919
 920     return 0;
 921 }
 922
 923 static int cbs_h2645_write_slice_data(CodedBitstreamContext *ctx,
 924                                       PutBitContext *pbc, const uint8_t *data,
 925                                       size_t data_size, int data_bit_start)
 926 {
 927     size_t rest  = data_size - (data_bit_start + 7) / 8;
 928     const uint8_t *pos = data + data_bit_start / 8;
 929
 930     av_assert0(data_bit_start >= 0 &&
 931                data_size > data_bit_start / 8);
 932
 933     if (data_size * 8 + 8 > put_bits_left(pbc))
 934         return AVERROR(ENOSPC);
 935
 936     if (!rest)
 937         goto rbsp_stop_one_bit;
 938
 939     // First copy the remaining bits of the first byte
 940     // The above check ensures that we do not accidentally
 941     // copy beyond the rbsp_stop_one_bit.
 942     if (data_bit_start % 8)
 943         put_bits(pbc, 8 - data_bit_start % 8,
 944                  *pos++ & MAX_UINT_BITS(8 - data_bit_start % 8));
 945
 946     if (put_bits_count(pbc) % 8 == 0) {
 947         // If the writer is aligned at this point,
 948         // memcpy can be used to improve performance.
 949         // This happens normally for CABAC.
 950         flush_put_bits(pbc);
 951         memcpy(put_bits_ptr(pbc), pos, rest);
 952         skip_put_bytes(pbc, rest);
 953     } else {
 954         // If not, we have to copy manually.
 955         // rbsp_stop_one_bit forces us to special-case
 956         // the last byte.
 957         uint8_t temp;
 958         int i;
 959
 960         for (; rest > 4; rest -= 4, pos += 4)
 961             put_bits32(pbc, AV_RB32(pos));
 962
 963         for (; rest > 1; rest--, pos++)
 964             put_bits(pbc, 8, *pos);
 965
 966     rbsp_stop_one_bit:
 967         temp = rest ? *pos : *pos & MAX_UINT_BITS(8 - data_bit_start % 8);
 968
 969         av_assert0(temp);
 970         i = ff_ctz(*pos);
 971         temp = temp >> i;
 972         i = rest ? (8 - i) : (8 - i - data_bit_start % 8);
 973         put_bits(pbc, i, temp);
 974         if (put_bits_count(pbc) % 8)
 975             put_bits(pbc, 8 - put_bits_count(pbc) % 8, 0);
 976     }
 977
 978     return 0;
 979 }
 980
 981 static int cbs_h264_write_nal_unit(CodedBitstreamContext *ctx,
 982                                    CodedBitstreamUnit *unit,
 983                                    PutBitContext *pbc)
 984 {
 985     int err;
 986
 987     switch (unit->type) {
 988     case H264_NAL_SPS:
 989         {
 990             H264RawSPS *sps = unit->content;
 991
 992             err = cbs_h264_write_sps(ctx, pbc, sps);
 993             if (err < 0)
 994                 return err;
 995
 996             err = cbs_h264_replace_sps(ctx, unit);
 997             if (err < 0)
 998                 return err;
 999         }
1000         break;
1001
1002     case H264_NAL_SPS_EXT:
1003         {
1004             H264RawSPSExtension *sps_ext = unit->content;
1005
1006             err = cbs_h264_write_sps_extension(ctx, pbc, sps_ext);
1007             if (err < 0)
1008                 return err;
1009         }
1010         break;
1011
1012     case H264_NAL_PPS:
1013         {
1014             H264RawPPS *pps = unit->content;
1015
1016             err = cbs_h264_write_pps(ctx, pbc, pps);
1017             if (err < 0)
1018                 return err;
1019
1020             err = cbs_h264_replace_pps(ctx, unit);
1021             if (err < 0)
1022                 return err;
1023         }
1024         break;
1025
1026     case H264_NAL_SLICE:
1027     case H264_NAL_IDR_SLICE:
1028     case H264_NAL_AUXILIARY_SLICE:
1029         {
1030             H264RawSlice *slice = unit->content;
1031
1032             err = cbs_h264_write_slice_header(ctx, pbc, &slice->header);
1033             if (err < 0)
1034                 return err;
1035
1036             if (slice->data) {
1037                 err = cbs_h2645_write_slice_data(ctx, pbc, slice->data,
1038                                                  slice->data_size,
1039                                                  slice->data_bit_start);
1040                 if (err < 0)
1041                     return err;
1042             } else {
1043                 // No slice data - that was just the header.
1044                 // (Bitstream may be unaligned!)
1045             }
1046         }
1047         break;
1048
1049     case H264_NAL_AUD:
1050         {
1051             err = cbs_h264_write_aud(ctx, pbc, unit->content);
1052             if (err < 0)
1053                 return err;
1054         }
1055         break;
1056
1057     case H264_NAL_SEI:
1058         {
1059             err = cbs_h264_write_sei(ctx, pbc, unit->content);
1060             if (err < 0)
1061                 return err;
1062         }
1063         break;
1064
1065     case H264_NAL_FILLER_DATA:
1066         {
1067             err = cbs_h264_write_filler(ctx, pbc, unit->content);
1068             if (err < 0)
1069                 return err;
1070         }
1071         break;
1072
1073     case H264_NAL_END_SEQUENCE:
1074         {
1075             err = cbs_h264_write_end_of_sequence(ctx, pbc, unit->content);
1076             if (err < 0)
1077                 return err;
1078         }
1079         break;
1080
1081     case H264_NAL_END_STREAM:
1082         {
1083             err = cbs_h264_write_end_of_stream(ctx, pbc, unit->content);
1084             if (err < 0)
1085                 return err;
1086         }
1087         break;
1088
1089     default:
1090         av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for "
1091                "NAL unit type %"PRIu32".\n", unit->type);
1092         return AVERROR_PATCHWELCOME;
1093     }
1094
1095     return 0;
1096 }
1097
1098 static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
1099                                    CodedBitstreamUnit *unit,
1100                                    PutBitContext *pbc)
1101 {
1102     int err;
1103
1104     switch (unit->type) {
1105     case HEVC_NAL_VPS:
1106         {
1107             H265RawVPS *vps = unit->content;
1108
1109             err = cbs_h265_write_vps(ctx, pbc, vps);
1110             if (err < 0)
1111                 return err;
1112
1113             err = cbs_h265_replace_vps(ctx, unit);
1114             if (err < 0)
1115                 return err;
1116         }
1117         break;
1118
1119     case HEVC_NAL_SPS:
1120         {
1121             H265RawSPS *sps = unit->content;
1122
1123             err = cbs_h265_write_sps(ctx, pbc, sps);
1124             if (err < 0)
1125                 return err;
1126
1127             err = cbs_h265_replace_sps(ctx, unit);
1128             if (err < 0)
1129                 return err;
1130         }
1131         break;
1132
1133     case HEVC_NAL_PPS:
1134         {
1135             H265RawPPS *pps = unit->content;
1136
1137             err = cbs_h265_write_pps(ctx, pbc, pps);
1138             if (err < 0)
1139                 return err;
1140
1141             err = cbs_h265_replace_pps(ctx, unit);
1142             if (err < 0)
1143                 return err;
1144         }
1145         break;
1146
1147     case HEVC_NAL_TRAIL_N:
1148     case HEVC_NAL_TRAIL_R:
1149     case HEVC_NAL_TSA_N:
1150     case HEVC_NAL_TSA_R:
1151     case HEVC_NAL_STSA_N:
1152     case HEVC_NAL_STSA_R:
1153     case HEVC_NAL_RADL_N:
1154     case HEVC_NAL_RADL_R:
1155     case HEVC_NAL_RASL_N:
1156     case HEVC_NAL_RASL_R:
1157     case HEVC_NAL_BLA_W_LP:
1158     case HEVC_NAL_BLA_W_RADL:
1159     case HEVC_NAL_BLA_N_LP:
1160     case HEVC_NAL_IDR_W_RADL:
1161     case HEVC_NAL_IDR_N_LP:
1162     case HEVC_NAL_CRA_NUT:
1163         {
1164             H265RawSlice *slice = unit->content;
1165
1166             err = cbs_h265_write_slice_segment_header(ctx, pbc, &slice->header);
1167             if (err < 0)
1168                 return err;
1169
1170             if (slice->data) {
1171                 err = cbs_h2645_write_slice_data(ctx, pbc, slice->data,
1172                                                  slice->data_size,
1173                                                  slice->data_bit_start);
1174                 if (err < 0)
1175                     return err;
1176             } else {
1177                 // No slice data - that was just the header.
1178             }
1179         }
1180         break;
1181
1182     case HEVC_NAL_AUD:
1183         {
1184             err = cbs_h265_write_aud(ctx, pbc, unit->content);
1185             if (err < 0)
1186                 return err;
1187         }
1188         break;
1189
1190     case HEVC_NAL_SEI_PREFIX:
1191     case HEVC_NAL_SEI_SUFFIX:
1192         {
1193             err = cbs_h265_write_sei(ctx, pbc, unit->content,
1194                                      unit->type == HEVC_NAL_SEI_PREFIX);
1195
1196             if (err < 0)
1197                 return err;
1198         }
1199         break;
1200
1201     default:
1202         av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for "
1203                "NAL unit type %"PRIu32".\n", unit->type);
1204         return AVERROR_PATCHWELCOME;
1205     }
1206
1207     return 0;
1208 }
1209
1210 static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
1211                                        CodedBitstreamFragment *frag)
1212 {
1213     uint8_t *data;
1214     size_t max_size, dp, sp;
1215     int err, i, zero_run;
1216
1217     for (i = 0; i < frag->nb_units; i++) {
1218         // Data should already all have been written when we get here.
1219         av_assert0(frag->units[i].data);
1220     }
1221
1222     max_size = 0;
1223     for (i = 0; i < frag->nb_units; i++) {
1224         // Start code + content with worst-case emulation prevention.
1225         max_size += 4 + frag->units[i].data_size * 3 / 2;
1226     }
1227
1228     data = av_realloc(NULL, max_size + AV_INPUT_BUFFER_PADDING_SIZE);
1229     if (!data)
1230         return AVERROR(ENOMEM);
1231
1232     dp = 0;
1233     for (i = 0; i < frag->nb_units; i++) {
1234         CodedBitstreamUnit *unit = &frag->units[i];
1235
1236         if (unit->data_bit_padding > 0) {
1237             if (i < frag->nb_units - 1)
1238                 av_log(ctx->log_ctx, AV_LOG_WARNING, "Probably invalid "
1239                        "unaligned padding on non-final NAL unit.\n");
1240             else
1241                 frag->data_bit_padding = unit->data_bit_padding;
1242         }
1243
1244         if ((ctx->codec->codec_id == AV_CODEC_ID_H264 &&
1245              (unit->type == H264_NAL_SPS ||
1246               unit->type == H264_NAL_PPS)) ||
1247             (ctx->codec->codec_id == AV_CODEC_ID_HEVC &&
1248              (unit->type == HEVC_NAL_VPS ||
1249               unit->type == HEVC_NAL_SPS ||
1250               unit->type == HEVC_NAL_PPS)) ||
1251             i == 0 /* (Assume this is the start of an access unit.) */) {
1252             // zero_byte
1253             data[dp++] = 0;
1254         }
1255         // start_code_prefix_one_3bytes
1256         data[dp++] = 0;
1257         data[dp++] = 0;
1258         data[dp++] = 1;
1259
1260         zero_run = 0;
1261         for (sp = 0; sp < unit->data_size; sp++) {
1262             if (zero_run < 2) {
1263                 if (unit->data[sp] == 0)
1264                     ++zero_run;
1265                 else
1266                     zero_run = 0;
1267             } else {
1268                 if ((unit->data[sp] & ~3) == 0) {
1269                     // emulation_prevention_three_byte
1270                     data[dp++] = 3;
1271                 }
1272                 zero_run = unit->data[sp] == 0;
1273             }
1274             data[dp++] = unit->data[sp];
1275         }
1276     }
1277
1278     av_assert0(dp <= max_size);
1279     err = av_reallocp(&data, dp + AV_INPUT_BUFFER_PADDING_SIZE);
1280     if (err)
1281         return err;
1282     memset(data + dp, 0, AV_INPUT_BUFFER_PADDING_SIZE);
1283
1284     frag->data_ref = av_buffer_create(data, dp + AV_INPUT_BUFFER_PADDING_SIZE,
1285                                       NULL, NULL, 0);
1286     if (!frag->data_ref) {
1287         av_freep(&data);
1288         return AVERROR(ENOMEM);
1289     }
1290
1291     frag->data = data;
1292     frag->data_size = dp;
1293
1294     return 0;
1295 }
1296
1297 static void cbs_h264_flush(CodedBitstreamContext *ctx)
1298 {
1299     CodedBitstreamH264Context *h264 = ctx->priv_data;
1300
1301     for (int i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++) {
1302         av_buffer_unref(&h264->sps_ref[i]);
1303         h264->sps[i] = NULL;
1304     }
1305     for (int i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++) {
1306         av_buffer_unref(&h264->pps_ref[i]);
1307         h264->pps[i] = NULL;
1308     }
1309
1310     h264->active_sps = NULL;
1311     h264->active_pps = NULL;
1312     h264->last_slice_nal_unit_type = 0;
1313 }
1314
1315 static void cbs_h264_close(CodedBitstreamContext *ctx)
1316 {
1317     CodedBitstreamH264Context *h264 = ctx->priv_data;
1318     int i;
1319
1320     ff_h2645_packet_uninit(&h264->common.read_packet);
1321
1322     for (i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++)
1323         av_buffer_unref(&h264->sps_ref[i]);
1324     for (i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++)
1325         av_buffer_unref(&h264->pps_ref[i]);
1326 }
1327
1328 static void cbs_h265_flush(CodedBitstreamContext *ctx)
1329 {
1330     CodedBitstreamH265Context *h265 = ctx->priv_data;
1331
1332     for (int i = 0; i < FF_ARRAY_ELEMS(h265->vps); i++) {
1333         av_buffer_unref(&h265->vps_ref[i]);
1334         h265->vps[i] = NULL;
1335     }
1336     for (int i = 0; i < FF_ARRAY_ELEMS(h265->sps); i++) {
1337         av_buffer_unref(&h265->sps_ref[i]);
1338         h265->sps[i] = NULL;
1339     }
1340     for (int i = 0; i < FF_ARRAY_ELEMS(h265->pps); i++) {
1341         av_buffer_unref(&h265->pps_ref[i]);
1342         h265->pps[i] = NULL;
1343     }
1344
1345     h265->active_vps = NULL;
1346     h265->active_sps = NULL;
1347     h265->active_pps = NULL;
1348 }
1349
1350 static void cbs_h265_close(CodedBitstreamContext *ctx)
1351 {
1352     CodedBitstreamH265Context *h265 = ctx->priv_data;
1353     int i;
1354
1355     ff_h2645_packet_uninit(&h265->common.read_packet);
1356
1357     for (i = 0; i < FF_ARRAY_ELEMS(h265->vps); i++)
1358         av_buffer_unref(&h265->vps_ref[i]);
1359     for (i = 0; i < FF_ARRAY_ELEMS(h265->sps); i++)
1360         av_buffer_unref(&h265->sps_ref[i]);
1361     for (i = 0; i < FF_ARRAY_ELEMS(h265->pps); i++)
1362         av_buffer_unref(&h265->pps_ref[i]);
1363 }
1364
1365 static void cbs_h264_free_sei_payload(H264RawSEIPayload *payload)
1366 {
1367     switch (payload->payload_type) {
1368     case H264_SEI_TYPE_BUFFERING_PERIOD:
1369     case H264_SEI_TYPE_PIC_TIMING:
1370     case H264_SEI_TYPE_PAN_SCAN_RECT:
1371     case H264_SEI_TYPE_RECOVERY_POINT:
1372     case H264_SEI_TYPE_DISPLAY_ORIENTATION:
1373     case H264_SEI_TYPE_MASTERING_DISPLAY_COLOUR_VOLUME:
1374     case H264_SEI_TYPE_ALTERNATIVE_TRANSFER:
1375         break;
1376     case H264_SEI_TYPE_USER_DATA_REGISTERED:
1377         av_buffer_unref(&payload->payload.user_data_registered.data_ref);
1378         break;
1379     case H264_SEI_TYPE_USER_DATA_UNREGISTERED:
1380         av_buffer_unref(&payload->payload.user_data_unregistered.data_ref);
1381         break;
1382     default:
1383         av_buffer_unref(&payload->payload.other.data_ref);
1384         break;
1385     }
1386 }
1387
1388 static void cbs_h264_free_sei(void *opaque, uint8_t *content)
1389 {
1390     H264RawSEI *sei = (H264RawSEI*)content;
1391     int i;
1392     for (i = 0; i < sei->payload_count; i++)
1393         cbs_h264_free_sei_payload(&sei->payload[i]);
1394     av_freep(&content);
1395 }
1396
1397 static const CodedBitstreamUnitTypeDescriptor cbs_h264_unit_types[] = {
1398     CBS_UNIT_TYPE_POD(H264_NAL_SPS,     H264RawSPS),
1399     CBS_UNIT_TYPE_POD(H264_NAL_SPS_EXT, H264RawSPSExtension),
1400
1401     CBS_UNIT_TYPE_INTERNAL_REF(H264_NAL_PPS, H264RawPPS, slice_group_id),
1402
1403     {
1404         .nb_unit_types  = 3,
1405         .unit_types     = {
1406             H264_NAL_IDR_SLICE,
1407             H264_NAL_SLICE,
1408             H264_NAL_AUXILIARY_SLICE,
1409         },
1410         .content_type   = CBS_CONTENT_TYPE_INTERNAL_REFS,
1411         .content_size   = sizeof(H264RawSlice),
1412         .nb_ref_offsets = 1,
1413         .ref_offsets    = { offsetof(H264RawSlice, data) },
1414     },
1415
1416     CBS_UNIT_TYPE_POD(H264_NAL_AUD,          H264RawAUD),
1417     CBS_UNIT_TYPE_POD(H264_NAL_FILLER_DATA,  H264RawFiller),
1418     CBS_UNIT_TYPE_POD(H264_NAL_END_SEQUENCE, H264RawNALUnitHeader),
1419     CBS_UNIT_TYPE_POD(H264_NAL_END_STREAM,   H264RawNALUnitHeader),
1420
1421     CBS_UNIT_TYPE_COMPLEX(H264_NAL_SEI, H264RawSEI, &cbs_h264_free_sei),
1422
1423     CBS_UNIT_TYPE_END_OF_LIST
1424 };
1425
1426 static void cbs_h265_free_sei_payload(H265RawSEIPayload *payload)
1427 {
1428     switch (payload->payload_type) {
1429     case HEVC_SEI_TYPE_BUFFERING_PERIOD:
1430     case HEVC_SEI_TYPE_PICTURE_TIMING:
1431     case HEVC_SEI_TYPE_PAN_SCAN_RECT:
1432     case HEVC_SEI_TYPE_RECOVERY_POINT:
1433     case HEVC_SEI_TYPE_DISPLAY_ORIENTATION:
1434     case HEVC_SEI_TYPE_ACTIVE_PARAMETER_SETS:
1435     case HEVC_SEI_TYPE_DECODED_PICTURE_HASH:
1436     case HEVC_SEI_TYPE_TIME_CODE:
1437     case HEVC_SEI_TYPE_MASTERING_DISPLAY_INFO:
1438     case HEVC_SEI_TYPE_CONTENT_LIGHT_LEVEL_INFO:
1439     case HEVC_SEI_TYPE_ALTERNATIVE_TRANSFER_CHARACTERISTICS:
1440     case HEVC_SEI_TYPE_ALPHA_CHANNEL_INFO:
1441         break;
1442     case HEVC_SEI_TYPE_USER_DATA_REGISTERED_ITU_T_T35:
1443         av_buffer_unref(&payload->payload.user_data_registered.data_ref);
1444         break;
1445     case HEVC_SEI_TYPE_USER_DATA_UNREGISTERED:
1446         av_buffer_unref(&payload->payload.user_data_unregistered.data_ref);
1447         break;
1448     default:
1449         av_buffer_unref(&payload->payload.other.data_ref);
1450         break;
1451     }
1452     av_buffer_unref(&payload->extension_data.data_ref);
1453 }
1454
1455 static void cbs_h265_free_sei(void *opaque, uint8_t *content)
1456 {
1457     H265RawSEI *sei = (H265RawSEI*)content;
1458     int i;
1459     for (i = 0; i < sei->payload_count; i++)
1460         cbs_h265_free_sei_payload(&sei->payload[i]);
1461     av_freep(&content);
1462 }
1463
1464 static const CodedBitstreamUnitTypeDescriptor cbs_h265_unit_types[] = {
1465     CBS_UNIT_TYPE_INTERNAL_REF(HEVC_NAL_VPS, H265RawVPS, extension_data.data),
1466     CBS_UNIT_TYPE_INTERNAL_REF(HEVC_NAL_SPS, H265RawSPS, extension_data.data),
1467     CBS_UNIT_TYPE_INTERNAL_REF(HEVC_NAL_PPS, H265RawPPS, extension_data.data),
1468
1469     CBS_UNIT_TYPE_POD(HEVC_NAL_AUD, H265RawAUD),
1470
1471     {
1472         // Slices of non-IRAP pictures.
1473         .nb_unit_types         = CBS_UNIT_TYPE_RANGE,
1474         .unit_type_range_start = HEVC_NAL_TRAIL_N,
1475         .unit_type_range_end   = HEVC_NAL_RASL_R,
1476
1477         .content_type   = CBS_CONTENT_TYPE_INTERNAL_REFS,
1478         .content_size   = sizeof(H265RawSlice),
1479         .nb_ref_offsets = 1,
1480         .ref_offsets    = { offsetof(H265RawSlice, data) },
1481     },
1482
1483     {
1484         // Slices of IRAP pictures.
1485         .nb_unit_types         = CBS_UNIT_TYPE_RANGE,
1486         .unit_type_range_start = HEVC_NAL_BLA_W_LP,
1487         .unit_type_range_end   = HEVC_NAL_CRA_NUT,
1488
1489         .content_type   = CBS_CONTENT_TYPE_INTERNAL_REFS,
1490         .content_size   = sizeof(H265RawSlice),
1491         .nb_ref_offsets = 1,
1492         .ref_offsets    = { offsetof(H265RawSlice, data) },
1493     },
1494
1495     {
1496         .nb_unit_types  = 2,
1497         .unit_types     = {
1498             HEVC_NAL_SEI_PREFIX,
1499             HEVC_NAL_SEI_SUFFIX
1500         },
1501         .content_type   = CBS_CONTENT_TYPE_COMPLEX,
1502         .content_size   = sizeof(H265RawSEI),
1503         .content_free   = &cbs_h265_free_sei,
1504     },
1505
1506     CBS_UNIT_TYPE_END_OF_LIST
1507 };
1508
1509 const CodedBitstreamType ff_cbs_type_h264 = {
1510     .codec_id          = AV_CODEC_ID_H264,
1511
1512     .priv_data_size    = sizeof(CodedBitstreamH264Context),
1513
1514     .unit_types        = cbs_h264_unit_types,
1515
1516     .split_fragment    = &cbs_h2645_split_fragment,
1517     .read_unit         = &cbs_h264_read_nal_unit,
1518     .write_unit        = &cbs_h264_write_nal_unit,
1519     .assemble_fragment = &cbs_h2645_assemble_fragment,
1520
1521     .flush             = &cbs_h264_flush,
1522     .close             = &cbs_h264_close,
1523 };
1524
1525 const CodedBitstreamType ff_cbs_type_h265 = {
1526     .codec_id          = AV_CODEC_ID_HEVC,
1527
1528     .priv_data_size    = sizeof(CodedBitstreamH265Context),
1529
1530     .unit_types        = cbs_h265_unit_types,
1531
1532     .split_fragment    = &cbs_h2645_split_fragment,
1533     .read_unit         = &cbs_h265_read_nal_unit,
1534     .write_unit        = &cbs_h265_write_nal_unit,
1535     .assemble_fragment = &cbs_h2645_assemble_fragment,
1536
1537     .flush             = &cbs_h265_flush,
1538     .close             = &cbs_h265_close,
1539 };
1540
1541 int ff_cbs_h264_add_sei_message(CodedBitstreamFragment *au,
1542                                 H264RawSEIPayload *payload)
1543 {
1544     H264RawSEI *sei = NULL;
1545     int err, i;
1546
1547     // Find an existing SEI NAL unit to add to.
1548     for (i = 0; i < au->nb_units; i++) {
1549         if (au->units[i].type == H264_NAL_SEI) {
1550             sei = au->units[i].content;
1551             if (sei->payload_count < H264_MAX_SEI_PAYLOADS)
1552                 break;
1553
1554             sei = NULL;
1555         }
1556     }
1557
1558     if (!sei) {
1559         // Need to make a new SEI NAL unit.  Insert it before the first
1560         // slice data NAL unit; if no slice data, add at the end.
1561         AVBufferRef *sei_ref;
1562
1563         sei = av_mallocz(sizeof(*sei));
1564         if (!sei) {
1565             err = AVERROR(ENOMEM);
1566             goto fail;
1567         }
1568
1569         sei->nal_unit_header.nal_unit_type = H264_NAL_SEI;
1570         sei->nal_unit_header.nal_ref_idc   = 0;
1571
1572         sei_ref = av_buffer_create((uint8_t*)sei, sizeof(*sei),
1573                                    &cbs_h264_free_sei, NULL, 0);
1574         if (!sei_ref) {
1575             av_freep(&sei);
1576             err = AVERROR(ENOMEM);
1577             goto fail;
1578         }
1579
1580         for (i = 0; i < au->nb_units; i++) {
1581             if (au->units[i].type == H264_NAL_SLICE ||
1582                 au->units[i].type == H264_NAL_IDR_SLICE)
1583                 break;
1584         }
1585
1586         err = ff_cbs_insert_unit_content(au, i, H264_NAL_SEI,
1587                                          sei, sei_ref);
1588         av_buffer_unref(&sei_ref);
1589         if (err < 0)
1590             goto fail;
1591     }
1592
1593     memcpy(&sei->payload[sei->payload_count], payload, sizeof(*payload));
1594     ++sei->payload_count;
1595
1596     return 0;
1597 fail:
1598     cbs_h264_free_sei_payload(payload);
1599     return err;
1600 }
1601
1602 void ff_cbs_h264_delete_sei_message(CodedBitstreamFragment *au,
1603                                     CodedBitstreamUnit *nal,
1604                                     int position)
1605 {
1606     H264RawSEI *sei = nal->content;
1607
1608     av_assert0(nal->type == H264_NAL_SEI);
1609     av_assert0(position >= 0 && position < sei->payload_count);
1610
1611     if (position == 0 && sei->payload_count == 1) {
1612         // Deleting NAL unit entirely.
1613         int i;
1614
1615         for (i = 0; i < au->nb_units; i++) {
1616             if (&au->units[i] == nal)
1617                 break;
1618         }
1619
1620         ff_cbs_delete_unit(au, i);
1621     } else {
1622         cbs_h264_free_sei_payload(&sei->payload[position]);
1623
1624         --sei->payload_count;
1625         memmove(sei->payload + position,
1626                 sei->payload + position + 1,
1627                 (sei->payload_count - position) * sizeof(*sei->payload));
1628     }
1629 }