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_NAME2(rw, codec, name) cbs_ ## codec ## _ ## rw ## _ ## name
 257 #define FUNC_NAME1(rw, codec, name) FUNC_NAME2(rw, codec, name)
 258 #define FUNC_H264(name) FUNC_NAME1(READWRITE, h264, name)
 259 #define FUNC_H265(name) FUNC_NAME1(READWRITE, h265, name)
 260 #define FUNC_SEI(name)  FUNC_NAME1(READWRITE, sei,  name)
 261
 262 #define SUBSCRIPTS(subs, ...) (subs > 0 ? ((int[subs + 1]){ subs, __VA_ARGS__ }) : NULL)
 263
 264 #define u(width, name, range_min, range_max) \
 265         xu(width, name, current->name, range_min, range_max, 0, )
 266 #define ub(width, name) \
 267         xu(width, name, current->name, 0, MAX_UINT_BITS(width), 0, )
 268 #define flag(name) ub(1, name)
 269 #define ue(name, range_min, range_max) \
 270         xue(name, current->name, range_min, range_max, 0, )
 271 #define i(width, name, range_min, range_max) \
 272         xi(width, name, current->name, range_min, range_max, 0, )
 273 #define ib(width, name) \
 274         xi(width, name, current->name, MIN_INT_BITS(width), MAX_INT_BITS(width), 0, )
 275 #define se(name, range_min, range_max) \
 276         xse(name, current->name, range_min, range_max, 0, )
 277
 278 #define us(width, name, range_min, range_max, subs, ...) \
 279         xu(width, name, current->name, range_min, range_max, subs, __VA_ARGS__)
 280 #define ubs(width, name, subs, ...) \
 281         xu(width, name, current->name, 0, MAX_UINT_BITS(width), subs, __VA_ARGS__)
 282 #define flags(name, subs, ...) \
 283         xu(1, name, current->name, 0, 1, subs, __VA_ARGS__)
 284 #define ues(name, range_min, range_max, subs, ...) \
 285         xue(name, current->name, range_min, range_max, subs, __VA_ARGS__)
 286 #define is(width, name, range_min, range_max, subs, ...) \
 287         xi(width, name, current->name, range_min, range_max, subs, __VA_ARGS__)
 288 #define ibs(width, name, subs, ...) \
 289         xi(width, name, current->name, MIN_INT_BITS(width), MAX_INT_BITS(width), subs, __VA_ARGS__)
 290 #define ses(name, range_min, range_max, subs, ...) \
 291         xse(name, current->name, range_min, range_max, subs, __VA_ARGS__)
 292
 293 #define fixed(width, name, value) do { \
 294         av_unused uint32_t fixed_value = value; \
 295         xu(width, name, fixed_value, value, value, 0, ); \
 296     } while (0)
 297
 298
 299 #define READ
 300 #define READWRITE read
 301 #define RWContext GetBitContext
 302
 303 #define xu(width, name, var, range_min, range_max, subs, ...) do { \
 304         uint32_t value; \
 305         CHECK(ff_cbs_read_unsigned(ctx, rw, width, #name, \
 306                                    SUBSCRIPTS(subs, __VA_ARGS__), \
 307                                    &value, range_min, range_max)); \
 308         var = value; \
 309     } while (0)
 310 #define xue(name, var, range_min, range_max, subs, ...) do { \
 311         uint32_t value; \
 312         CHECK(cbs_read_ue_golomb(ctx, rw, #name, \
 313                                  SUBSCRIPTS(subs, __VA_ARGS__), \
 314                                  &value, range_min, range_max)); \
 315         var = value; \
 316     } while (0)
 317 #define xi(width, name, var, range_min, range_max, subs, ...) do { \
 318         int32_t value; \
 319         CHECK(ff_cbs_read_signed(ctx, rw, width, #name, \
 320                                  SUBSCRIPTS(subs, __VA_ARGS__), \
 321                                  &value, range_min, range_max)); \
 322         var = value; \
 323     } while (0)
 324 #define xse(name, var, range_min, range_max, subs, ...) do { \
 325         int32_t value; \
 326         CHECK(cbs_read_se_golomb(ctx, rw, #name, \
 327                                  SUBSCRIPTS(subs, __VA_ARGS__), \
 328                                  &value, range_min, range_max)); \
 329         var = value; \
 330     } while (0)
 331
 332
 333 #define infer(name, value) do { \
 334         current->name = value; \
 335     } while (0)
 336
 337 static int cbs_h2645_read_more_rbsp_data(GetBitContext *gbc)
 338 {
 339     int bits_left = get_bits_left(gbc);
 340     if (bits_left > 8)
 341         return 1;
 342     if (bits_left == 0)
 343         return 0;
 344     if (show_bits(gbc, bits_left) & MAX_UINT_BITS(bits_left - 1))
 345         return 1;
 346     return 0;
 347 }
 348
 349 #define more_rbsp_data(var) ((var) = cbs_h2645_read_more_rbsp_data(rw))
 350
 351 #define bit_position(rw)   (get_bits_count(rw))
 352 #define byte_alignment(rw) (get_bits_count(rw) % 8)
 353
 354 #define allocate(name, size) do { \
 355         name ## _ref = av_buffer_allocz(size + \
 356                                         AV_INPUT_BUFFER_PADDING_SIZE); \
 357         if (!name ## _ref) \
 358             return AVERROR(ENOMEM); \
 359         name = name ## _ref->data; \
 360     } while (0)
 361
 362 #define FUNC(name) FUNC_SEI(name)
 363 #include "cbs_sei_syntax_template.c"
 364 #undef FUNC
 365
 366 #define FUNC(name) FUNC_H264(name)
 367 #include "cbs_h264_syntax_template.c"
 368 #undef FUNC
 369
 370 #define FUNC(name) FUNC_H265(name)
 371 #include "cbs_h265_syntax_template.c"
 372 #undef FUNC
 373
 374 #undef READ
 375 #undef READWRITE
 376 #undef RWContext
 377 #undef xu
 378 #undef xi
 379 #undef xue
 380 #undef xse
 381 #undef infer
 382 #undef more_rbsp_data
 383 #undef bit_position
 384 #undef byte_alignment
 385 #undef allocate
 386
 387
 388 #define WRITE
 389 #define READWRITE write
 390 #define RWContext PutBitContext
 391
 392 #define xu(width, name, var, range_min, range_max, subs, ...) do { \
 393         uint32_t value = var; \
 394         CHECK(ff_cbs_write_unsigned(ctx, rw, width, #name, \
 395                                     SUBSCRIPTS(subs, __VA_ARGS__), \
 396                                     value, range_min, range_max)); \
 397     } while (0)
 398 #define xue(name, var, range_min, range_max, subs, ...) do { \
 399         uint32_t value = var; \
 400         CHECK(cbs_write_ue_golomb(ctx, rw, #name, \
 401                                   SUBSCRIPTS(subs, __VA_ARGS__), \
 402                                   value, range_min, range_max)); \
 403     } while (0)
 404 #define xi(width, name, var, range_min, range_max, subs, ...) do { \
 405         int32_t value = var; \
 406         CHECK(ff_cbs_write_signed(ctx, rw, width, #name, \
 407                                   SUBSCRIPTS(subs, __VA_ARGS__), \
 408                                   value, range_min, range_max)); \
 409     } while (0)
 410 #define xse(name, var, range_min, range_max, subs, ...) do { \
 411         int32_t value = var; \
 412         CHECK(cbs_write_se_golomb(ctx, rw, #name, \
 413                                   SUBSCRIPTS(subs, __VA_ARGS__), \
 414                                   value, range_min, range_max)); \
 415     } while (0)
 416
 417 #define infer(name, value) do { \
 418         if (current->name != (value)) { \
 419             av_log(ctx->log_ctx, AV_LOG_ERROR, \
 420                    "%s does not match inferred value: " \
 421                    "%"PRId64", but should be %"PRId64".\n", \
 422                    #name, (int64_t)current->name, (int64_t)(value)); \
 423             return AVERROR_INVALIDDATA; \
 424         } \
 425     } while (0)
 426
 427 #define more_rbsp_data(var) (var)
 428
 429 #define bit_position(rw)   (put_bits_count(rw))
 430 #define byte_alignment(rw) (put_bits_count(rw) % 8)
 431
 432 #define allocate(name, size) do { \
 433         if (!name) { \
 434             av_log(ctx->log_ctx, AV_LOG_ERROR, "%s must be set " \
 435                    "for writing.\n", #name); \
 436             return AVERROR_INVALIDDATA; \
 437         } \
 438     } while (0)
 439
 440 #define FUNC(name) FUNC_SEI(name)
 441 #include "cbs_sei_syntax_template.c"
 442 #undef FUNC
 443
 444 #define FUNC(name) FUNC_H264(name)
 445 #include "cbs_h264_syntax_template.c"
 446 #undef FUNC
 447
 448 #define FUNC(name) FUNC_H265(name)
 449 #include "cbs_h265_syntax_template.c"
 450 #undef FUNC
 451
 452 #undef WRITE
 453 #undef READWRITE
 454 #undef RWContext
 455 #undef xu
 456 #undef xi
 457 #undef xue
 458 #undef xse
 459 #undef u
 460 #undef i
 461 #undef flag
 462 #undef ue
 463 #undef se
 464 #undef infer
 465 #undef more_rbsp_data
 466 #undef bit_position
 467 #undef byte_alignment
 468 #undef allocate
 469
 470
 471 static int cbs_h2645_fragment_add_nals(CodedBitstreamContext *ctx,
 472                                        CodedBitstreamFragment *frag,
 473                                        const H2645Packet *packet)
 474 {
 475     int err, i;
 476
 477     for (i = 0; i < packet->nb_nals; i++) {
 478         const H2645NAL *nal = &packet->nals[i];
 479         AVBufferRef *ref;
 480         size_t size = nal->size;
 481
 482         if (nal->nuh_layer_id > 0)
 483             continue;
 484
 485         // Remove trailing zeroes.
 486         while (size > 0 && nal->data[size - 1] == 0)
 487             --size;
 488         if (size == 0) {
 489             av_log(ctx->log_ctx, AV_LOG_VERBOSE, "Discarding empty 0 NAL unit\n");
 490             continue;
 491         }
 492
 493         ref = (nal->data == nal->raw_data) ? frag->data_ref
 494                                            : packet->rbsp.rbsp_buffer_ref;
 495
 496         err = ff_cbs_insert_unit_data(frag, -1, nal->type,
 497                             (uint8_t*)nal->data, size, ref);
 498         if (err < 0)
 499             return err;
 500     }
 501
 502     return 0;
 503 }
 504
 505 static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
 506                                     CodedBitstreamFragment *frag,
 507                                     int header)
 508 {
 509     enum AVCodecID codec_id = ctx->codec->codec_id;
 510     CodedBitstreamH2645Context *priv = ctx->priv_data;
 511     GetByteContext gbc;
 512     int err;
 513
 514     av_assert0(frag->data && frag->nb_units == 0);
 515     if (frag->data_size == 0)
 516         return 0;
 517
 518     if (header && frag->data[0] && codec_id == AV_CODEC_ID_H264) {
 519         // AVCC header.
 520         size_t size, start, end;
 521         int i, count, version;
 522
 523         priv->mp4 = 1;
 524
 525         bytestream2_init(&gbc, frag->data, frag->data_size);
 526
 527         if (bytestream2_get_bytes_left(&gbc) < 6)
 528             return AVERROR_INVALIDDATA;
 529
 530         version = bytestream2_get_byte(&gbc);
 531         if (version != 1) {
 532             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid AVCC header: "
 533                    "first byte %u.\n", version);
 534             return AVERROR_INVALIDDATA;
 535         }
 536
 537         bytestream2_skip(&gbc, 3);
 538         priv->nal_length_size = (bytestream2_get_byte(&gbc) & 3) + 1;
 539
 540         // SPS array.
 541         count = bytestream2_get_byte(&gbc) & 0x1f;
 542         start = bytestream2_tell(&gbc);
 543         for (i = 0; i < count; i++) {
 544             if (bytestream2_get_bytes_left(&gbc) < 2 * (count - i))
 545                 return AVERROR_INVALIDDATA;
 546             size = bytestream2_get_be16(&gbc);
 547             if (bytestream2_get_bytes_left(&gbc) < size)
 548                 return AVERROR_INVALIDDATA;
 549             bytestream2_skip(&gbc, size);
 550         }
 551         end = bytestream2_tell(&gbc);
 552
 553         err = ff_h2645_packet_split(&priv->read_packet,
 554                                     frag->data + start, end - start,
 555                                     ctx->log_ctx, 1, 2, AV_CODEC_ID_H264, 1, 1);
 556         if (err < 0) {
 557             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC SPS array.\n");
 558             return err;
 559         }
 560         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 561         if (err < 0)
 562             return err;
 563
 564         // PPS array.
 565         count = bytestream2_get_byte(&gbc);
 566         start = bytestream2_tell(&gbc);
 567         for (i = 0; i < count; i++) {
 568             if (bytestream2_get_bytes_left(&gbc) < 2 * (count - i))
 569                 return AVERROR_INVALIDDATA;
 570             size = bytestream2_get_be16(&gbc);
 571             if (bytestream2_get_bytes_left(&gbc) < size)
 572                 return AVERROR_INVALIDDATA;
 573             bytestream2_skip(&gbc, size);
 574         }
 575         end = bytestream2_tell(&gbc);
 576
 577         err = ff_h2645_packet_split(&priv->read_packet,
 578                                     frag->data + start, end - start,
 579                                     ctx->log_ctx, 1, 2, AV_CODEC_ID_H264, 1, 1);
 580         if (err < 0) {
 581             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC PPS array.\n");
 582             return err;
 583         }
 584         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 585         if (err < 0)
 586             return err;
 587
 588         if (bytestream2_get_bytes_left(&gbc) > 0) {
 589             av_log(ctx->log_ctx, AV_LOG_WARNING, "%u bytes left at end of AVCC "
 590                    "header.\n", bytestream2_get_bytes_left(&gbc));
 591         }
 592
 593     } else if (header && frag->data[0] && codec_id == AV_CODEC_ID_HEVC) {
 594         // HVCC header.
 595         size_t size, start, end;
 596         int i, j, nb_arrays, nal_unit_type, nb_nals, version;
 597
 598         priv->mp4 = 1;
 599
 600         bytestream2_init(&gbc, frag->data, frag->data_size);
 601
 602         if (bytestream2_get_bytes_left(&gbc) < 23)
 603             return AVERROR_INVALIDDATA;
 604
 605         version = bytestream2_get_byte(&gbc);
 606         if (version != 1) {
 607             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid HVCC header: "
 608                    "first byte %u.\n", version);
 609             return AVERROR_INVALIDDATA;
 610         }
 611
 612         bytestream2_skip(&gbc, 20);
 613         priv->nal_length_size = (bytestream2_get_byte(&gbc) & 3) + 1;
 614
 615         nb_arrays = bytestream2_get_byte(&gbc);
 616         for (i = 0; i < nb_arrays; i++) {
 617             nal_unit_type = bytestream2_get_byte(&gbc) & 0x3f;
 618             nb_nals = bytestream2_get_be16(&gbc);
 619
 620             start = bytestream2_tell(&gbc);
 621             for (j = 0; j < nb_nals; j++) {
 622                 if (bytestream2_get_bytes_left(&gbc) < 2)
 623                     return AVERROR_INVALIDDATA;
 624                 size = bytestream2_get_be16(&gbc);
 625                 if (bytestream2_get_bytes_left(&gbc) < size)
 626                     return AVERROR_INVALIDDATA;
 627                 bytestream2_skip(&gbc, size);
 628             }
 629             end = bytestream2_tell(&gbc);
 630
 631             err = ff_h2645_packet_split(&priv->read_packet,
 632                                         frag->data + start, end - start,
 633                                         ctx->log_ctx, 1, 2, AV_CODEC_ID_HEVC, 1, 1);
 634             if (err < 0) {
 635                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split "
 636                        "HVCC array %d (%d NAL units of type %d).\n",
 637                        i, nb_nals, nal_unit_type);
 638                 return err;
 639             }
 640             err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 641             if (err < 0)
 642                 return err;
 643         }
 644
 645     } else {
 646         // Annex B, or later MP4 with already-known parameters.
 647
 648         err = ff_h2645_packet_split(&priv->read_packet,
 649                                     frag->data, frag->data_size,
 650                                     ctx->log_ctx,
 651                                     priv->mp4, priv->nal_length_size,
 652                                     codec_id, 1, 1);
 653         if (err < 0)
 654             return err;
 655
 656         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 657         if (err < 0)
 658             return err;
 659     }
 660
 661     return 0;
 662 }
 663
 664 #define cbs_h2645_replace_ps(h26n, ps_name, ps_var, id_element) \
 665 static int cbs_h26 ## h26n ## _replace_ ## ps_var(CodedBitstreamContext *ctx, \
 666                                                   CodedBitstreamUnit *unit)  \
 667 { \
 668     CodedBitstreamH26 ## h26n ## Context *priv = ctx->priv_data; \
 669     H26 ## h26n ## Raw ## ps_name *ps_var = unit->content; \
 670     unsigned int id = ps_var->id_element; \
 671     int err; \
 672     if (id >= FF_ARRAY_ELEMS(priv->ps_var)) { \
 673         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid " #ps_name \
 674                " id : %d.\n", id); \
 675         return AVERROR_INVALIDDATA; \
 676     } \
 677     err = ff_cbs_make_unit_refcounted(ctx, unit); \
 678     if (err < 0) \
 679         return err; \
 680     if (priv->ps_var[id] == priv->active_ ## ps_var) \
 681         priv->active_ ## ps_var = NULL ; \
 682     av_buffer_unref(&priv->ps_var ## _ref[id]); \
 683     av_assert0(unit->content_ref); \
 684     priv->ps_var ## _ref[id] = av_buffer_ref(unit->content_ref); \
 685     if (!priv->ps_var ## _ref[id]) \
 686         return AVERROR(ENOMEM); \
 687     priv->ps_var[id] = (H26 ## h26n ## Raw ## ps_name *)priv->ps_var ## _ref[id]->data; \
 688     return 0; \
 689 }
 690
 691 cbs_h2645_replace_ps(4, SPS, sps, seq_parameter_set_id)
 692 cbs_h2645_replace_ps(4, PPS, pps, pic_parameter_set_id)
 693 cbs_h2645_replace_ps(5, VPS, vps, vps_video_parameter_set_id)
 694 cbs_h2645_replace_ps(5, SPS, sps, sps_seq_parameter_set_id)
 695 cbs_h2645_replace_ps(5, PPS, pps, pps_pic_parameter_set_id)
 696
 697 static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
 698                                   CodedBitstreamUnit *unit)
 699 {
 700     GetBitContext gbc;
 701     int err;
 702
 703     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
 704     if (err < 0)
 705         return err;
 706
 707     err = ff_cbs_alloc_unit_content2(ctx, unit);
 708     if (err < 0)
 709         return err;
 710
 711     switch (unit->type) {
 712     case H264_NAL_SPS:
 713         {
 714             H264RawSPS *sps = unit->content;
 715
 716             err = cbs_h264_read_sps(ctx, &gbc, sps);
 717             if (err < 0)
 718                 return err;
 719
 720             err = cbs_h264_replace_sps(ctx, unit);
 721             if (err < 0)
 722                 return err;
 723         }
 724         break;
 725
 726     case H264_NAL_SPS_EXT:
 727         {
 728             err = cbs_h264_read_sps_extension(ctx, &gbc, unit->content);
 729             if (err < 0)
 730                 return err;
 731         }
 732         break;
 733
 734     case H264_NAL_PPS:
 735         {
 736             H264RawPPS *pps = unit->content;
 737
 738             err = cbs_h264_read_pps(ctx, &gbc, pps);
 739             if (err < 0)
 740                 return err;
 741
 742             err = cbs_h264_replace_pps(ctx, unit);
 743             if (err < 0)
 744                 return err;
 745         }
 746         break;
 747
 748     case H264_NAL_SLICE:
 749     case H264_NAL_IDR_SLICE:
 750     case H264_NAL_AUXILIARY_SLICE:
 751         {
 752             H264RawSlice *slice = unit->content;
 753             int pos, len;
 754
 755             err = cbs_h264_read_slice_header(ctx, &gbc, &slice->header);
 756             if (err < 0)
 757                 return err;
 758
 759             if (!cbs_h2645_read_more_rbsp_data(&gbc))
 760                 return AVERROR_INVALIDDATA;
 761
 762             pos = get_bits_count(&gbc);
 763             len = unit->data_size;
 764
 765             slice->data_size = len - pos / 8;
 766             slice->data_ref  = av_buffer_ref(unit->data_ref);
 767             if (!slice->data_ref)
 768                 return AVERROR(ENOMEM);
 769             slice->data = unit->data + pos / 8;
 770             slice->data_bit_start = pos % 8;
 771         }
 772         break;
 773
 774     case H264_NAL_AUD:
 775         {
 776             err = cbs_h264_read_aud(ctx, &gbc, unit->content);
 777             if (err < 0)
 778                 return err;
 779         }
 780         break;
 781
 782     case H264_NAL_SEI:
 783         {
 784             err = cbs_h264_read_sei(ctx, &gbc, unit->content);
 785             if (err < 0)
 786                 return err;
 787         }
 788         break;
 789
 790     case H264_NAL_FILLER_DATA:
 791         {
 792             err = cbs_h264_read_filler(ctx, &gbc, unit->content);
 793             if (err < 0)
 794                 return err;
 795         }
 796         break;
 797
 798     case H264_NAL_END_SEQUENCE:
 799     case H264_NAL_END_STREAM:
 800         {
 801             err = (unit->type == H264_NAL_END_SEQUENCE ?
 802                    cbs_h264_read_end_of_sequence :
 803                    cbs_h264_read_end_of_stream)(ctx, &gbc, unit->content);
 804             if (err < 0)
 805                 return err;
 806         }
 807         break;
 808
 809     default:
 810         return AVERROR(ENOSYS);
 811     }
 812
 813     return 0;
 814 }
 815
 816 static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
 817                                   CodedBitstreamUnit *unit)
 818 {
 819     GetBitContext gbc;
 820     int err;
 821
 822     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
 823     if (err < 0)
 824         return err;
 825
 826     err = ff_cbs_alloc_unit_content2(ctx, unit);
 827     if (err < 0)
 828         return err;
 829
 830     switch (unit->type) {
 831     case HEVC_NAL_VPS:
 832         {
 833             H265RawVPS *vps = unit->content;
 834
 835             err = cbs_h265_read_vps(ctx, &gbc, vps);
 836             if (err < 0)
 837                 return err;
 838
 839             err = cbs_h265_replace_vps(ctx, unit);
 840             if (err < 0)
 841                 return err;
 842         }
 843         break;
 844     case HEVC_NAL_SPS:
 845         {
 846             H265RawSPS *sps = unit->content;
 847
 848             err = cbs_h265_read_sps(ctx, &gbc, sps);
 849             if (err < 0)
 850                 return err;
 851
 852             err = cbs_h265_replace_sps(ctx, unit);
 853             if (err < 0)
 854                 return err;
 855         }
 856         break;
 857
 858     case HEVC_NAL_PPS:
 859         {
 860             H265RawPPS *pps = unit->content;
 861
 862             err = cbs_h265_read_pps(ctx, &gbc, pps);
 863             if (err < 0)
 864                 return err;
 865
 866             err = cbs_h265_replace_pps(ctx, unit);
 867             if (err < 0)
 868                 return err;
 869         }
 870         break;
 871
 872     case HEVC_NAL_TRAIL_N:
 873     case HEVC_NAL_TRAIL_R:
 874     case HEVC_NAL_TSA_N:
 875     case HEVC_NAL_TSA_R:
 876     case HEVC_NAL_STSA_N:
 877     case HEVC_NAL_STSA_R:
 878     case HEVC_NAL_RADL_N:
 879     case HEVC_NAL_RADL_R:
 880     case HEVC_NAL_RASL_N:
 881     case HEVC_NAL_RASL_R:
 882     case HEVC_NAL_BLA_W_LP:
 883     case HEVC_NAL_BLA_W_RADL:
 884     case HEVC_NAL_BLA_N_LP:
 885     case HEVC_NAL_IDR_W_RADL:
 886     case HEVC_NAL_IDR_N_LP:
 887     case HEVC_NAL_CRA_NUT:
 888         {
 889             H265RawSlice *slice = unit->content;
 890             int pos, len;
 891
 892             err = cbs_h265_read_slice_segment_header(ctx, &gbc, &slice->header);
 893             if (err < 0)
 894                 return err;
 895
 896             if (!cbs_h2645_read_more_rbsp_data(&gbc))
 897                 return AVERROR_INVALIDDATA;
 898
 899             pos = get_bits_count(&gbc);
 900             len = unit->data_size;
 901
 902             slice->data_size = len - pos / 8;
 903             slice->data_ref  = av_buffer_ref(unit->data_ref);
 904             if (!slice->data_ref)
 905                 return AVERROR(ENOMEM);
 906             slice->data = unit->data + pos / 8;
 907             slice->data_bit_start = pos % 8;
 908         }
 909         break;
 910
 911     case HEVC_NAL_AUD:
 912         {
 913             err = cbs_h265_read_aud(ctx, &gbc, unit->content);
 914             if (err < 0)
 915                 return err;
 916         }
 917         break;
 918
 919     case HEVC_NAL_SEI_PREFIX:
 920     case HEVC_NAL_SEI_SUFFIX:
 921         {
 922             err = cbs_h265_read_sei(ctx, &gbc, unit->content,
 923                                     unit->type == HEVC_NAL_SEI_PREFIX);
 924
 925             if (err < 0)
 926                 return err;
 927         }
 928         break;
 929
 930     default:
 931         return AVERROR(ENOSYS);
 932     }
 933
 934     return 0;
 935 }
 936
 937 static int cbs_h2645_write_slice_data(CodedBitstreamContext *ctx,
 938                                       PutBitContext *pbc, const uint8_t *data,
 939                                       size_t data_size, int data_bit_start)
 940 {
 941     size_t rest  = data_size - (data_bit_start + 7) / 8;
 942     const uint8_t *pos = data + data_bit_start / 8;
 943
 944     av_assert0(data_bit_start >= 0 &&
 945                data_size > data_bit_start / 8);
 946
 947     if (data_size * 8 + 8 > put_bits_left(pbc))
 948         return AVERROR(ENOSPC);
 949
 950     if (!rest)
 951         goto rbsp_stop_one_bit;
 952
 953     // First copy the remaining bits of the first byte
 954     // The above check ensures that we do not accidentally
 955     // copy beyond the rbsp_stop_one_bit.
 956     if (data_bit_start % 8)
 957         put_bits(pbc, 8 - data_bit_start % 8,
 958                  *pos++ & MAX_UINT_BITS(8 - data_bit_start % 8));
 959
 960     if (put_bits_count(pbc) % 8 == 0) {
 961         // If the writer is aligned at this point,
 962         // memcpy can be used to improve performance.
 963         // This happens normally for CABAC.
 964         flush_put_bits(pbc);
 965         memcpy(put_bits_ptr(pbc), pos, rest);
 966         skip_put_bytes(pbc, rest);
 967     } else {
 968         // If not, we have to copy manually.
 969         // rbsp_stop_one_bit forces us to special-case
 970         // the last byte.
 971         uint8_t temp;
 972         int i;
 973
 974         for (; rest > 4; rest -= 4, pos += 4)
 975             put_bits32(pbc, AV_RB32(pos));
 976
 977         for (; rest > 1; rest--, pos++)
 978             put_bits(pbc, 8, *pos);
 979
 980     rbsp_stop_one_bit:
 981         temp = rest ? *pos : *pos & MAX_UINT_BITS(8 - data_bit_start % 8);
 982
 983         av_assert0(temp);
 984         i = ff_ctz(*pos);
 985         temp = temp >> i;
 986         i = rest ? (8 - i) : (8 - i - data_bit_start % 8);
 987         put_bits(pbc, i, temp);
 988         if (put_bits_count(pbc) % 8)
 989             put_bits(pbc, 8 - put_bits_count(pbc) % 8, 0);
 990     }
 991
 992     return 0;
 993 }
 994
 995 static int cbs_h264_write_nal_unit(CodedBitstreamContext *ctx,
 996                                    CodedBitstreamUnit *unit,
 997                                    PutBitContext *pbc)
 998 {
 999     int err;
1000
1001     switch (unit->type) {
1002     case H264_NAL_SPS:
1003         {
1004             H264RawSPS *sps = unit->content;
1005
1006             err = cbs_h264_write_sps(ctx, pbc, sps);
1007             if (err < 0)
1008                 return err;
1009
1010             err = cbs_h264_replace_sps(ctx, unit);
1011             if (err < 0)
1012                 return err;
1013         }
1014         break;
1015
1016     case H264_NAL_SPS_EXT:
1017         {
1018             H264RawSPSExtension *sps_ext = unit->content;
1019
1020             err = cbs_h264_write_sps_extension(ctx, pbc, sps_ext);
1021             if (err < 0)
1022                 return err;
1023         }
1024         break;
1025
1026     case H264_NAL_PPS:
1027         {
1028             H264RawPPS *pps = unit->content;
1029
1030             err = cbs_h264_write_pps(ctx, pbc, pps);
1031             if (err < 0)
1032                 return err;
1033
1034             err = cbs_h264_replace_pps(ctx, unit);
1035             if (err < 0)
1036                 return err;
1037         }
1038         break;
1039
1040     case H264_NAL_SLICE:
1041     case H264_NAL_IDR_SLICE:
1042     case H264_NAL_AUXILIARY_SLICE:
1043         {
1044             H264RawSlice *slice = unit->content;
1045
1046             err = cbs_h264_write_slice_header(ctx, pbc, &slice->header);
1047             if (err < 0)
1048                 return err;
1049
1050             if (slice->data) {
1051                 err = cbs_h2645_write_slice_data(ctx, pbc, slice->data,
1052                                                  slice->data_size,
1053                                                  slice->data_bit_start);
1054                 if (err < 0)
1055                     return err;
1056             } else {
1057                 // No slice data - that was just the header.
1058                 // (Bitstream may be unaligned!)
1059             }
1060         }
1061         break;
1062
1063     case H264_NAL_AUD:
1064         {
1065             err = cbs_h264_write_aud(ctx, pbc, unit->content);
1066             if (err < 0)
1067                 return err;
1068         }
1069         break;
1070
1071     case H264_NAL_SEI:
1072         {
1073             err = cbs_h264_write_sei(ctx, pbc, unit->content);
1074             if (err < 0)
1075                 return err;
1076         }
1077         break;
1078
1079     case H264_NAL_FILLER_DATA:
1080         {
1081             err = cbs_h264_write_filler(ctx, pbc, unit->content);
1082             if (err < 0)
1083                 return err;
1084         }
1085         break;
1086
1087     case H264_NAL_END_SEQUENCE:
1088         {
1089             err = cbs_h264_write_end_of_sequence(ctx, pbc, unit->content);
1090             if (err < 0)
1091                 return err;
1092         }
1093         break;
1094
1095     case H264_NAL_END_STREAM:
1096         {
1097             err = cbs_h264_write_end_of_stream(ctx, pbc, unit->content);
1098             if (err < 0)
1099                 return err;
1100         }
1101         break;
1102
1103     default:
1104         av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for "
1105                "NAL unit type %"PRIu32".\n", unit->type);
1106         return AVERROR_PATCHWELCOME;
1107     }
1108
1109     return 0;
1110 }
1111
1112 static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
1113                                    CodedBitstreamUnit *unit,
1114                                    PutBitContext *pbc)
1115 {
1116     int err;
1117
1118     switch (unit->type) {
1119     case HEVC_NAL_VPS:
1120         {
1121             H265RawVPS *vps = unit->content;
1122
1123             err = cbs_h265_write_vps(ctx, pbc, vps);
1124             if (err < 0)
1125                 return err;
1126
1127             err = cbs_h265_replace_vps(ctx, unit);
1128             if (err < 0)
1129                 return err;
1130         }
1131         break;
1132
1133     case HEVC_NAL_SPS:
1134         {
1135             H265RawSPS *sps = unit->content;
1136
1137             err = cbs_h265_write_sps(ctx, pbc, sps);
1138             if (err < 0)
1139                 return err;
1140
1141             err = cbs_h265_replace_sps(ctx, unit);
1142             if (err < 0)
1143                 return err;
1144         }
1145         break;
1146
1147     case HEVC_NAL_PPS:
1148         {
1149             H265RawPPS *pps = unit->content;
1150
1151             err = cbs_h265_write_pps(ctx, pbc, pps);
1152             if (err < 0)
1153                 return err;
1154
1155             err = cbs_h265_replace_pps(ctx, unit);
1156             if (err < 0)
1157                 return err;
1158         }
1159         break;
1160
1161     case HEVC_NAL_TRAIL_N:
1162     case HEVC_NAL_TRAIL_R:
1163     case HEVC_NAL_TSA_N:
1164     case HEVC_NAL_TSA_R:
1165     case HEVC_NAL_STSA_N:
1166     case HEVC_NAL_STSA_R:
1167     case HEVC_NAL_RADL_N:
1168     case HEVC_NAL_RADL_R:
1169     case HEVC_NAL_RASL_N:
1170     case HEVC_NAL_RASL_R:
1171     case HEVC_NAL_BLA_W_LP:
1172     case HEVC_NAL_BLA_W_RADL:
1173     case HEVC_NAL_BLA_N_LP:
1174     case HEVC_NAL_IDR_W_RADL:
1175     case HEVC_NAL_IDR_N_LP:
1176     case HEVC_NAL_CRA_NUT:
1177         {
1178             H265RawSlice *slice = unit->content;
1179
1180             err = cbs_h265_write_slice_segment_header(ctx, pbc, &slice->header);
1181             if (err < 0)
1182                 return err;
1183
1184             if (slice->data) {
1185                 err = cbs_h2645_write_slice_data(ctx, pbc, slice->data,
1186                                                  slice->data_size,
1187                                                  slice->data_bit_start);
1188                 if (err < 0)
1189                     return err;
1190             } else {
1191                 // No slice data - that was just the header.
1192             }
1193         }
1194         break;
1195
1196     case HEVC_NAL_AUD:
1197         {
1198             err = cbs_h265_write_aud(ctx, pbc, unit->content);
1199             if (err < 0)
1200                 return err;
1201         }
1202         break;
1203
1204     case HEVC_NAL_SEI_PREFIX:
1205     case HEVC_NAL_SEI_SUFFIX:
1206         {
1207             err = cbs_h265_write_sei(ctx, pbc, unit->content,
1208                                      unit->type == HEVC_NAL_SEI_PREFIX);
1209
1210             if (err < 0)
1211                 return err;
1212         }
1213         break;
1214
1215     default:
1216         av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for "
1217                "NAL unit type %"PRIu32".\n", unit->type);
1218         return AVERROR_PATCHWELCOME;
1219     }
1220
1221     return 0;
1222 }
1223
1224 static int cbs_h2645_unit_requires_zero_byte(enum AVCodecID codec_id,
1225                                              CodedBitstreamUnitType type,
1226                                              int nal_unit_index)
1227 {
1228     // Section B.1.2 in H.264, section B.2.2 in H.265.
1229     if (nal_unit_index == 0) {
1230         // Assume that this is the first NAL unit in an access unit.
1231         return 1;
1232     }
1233     if (codec_id == AV_CODEC_ID_H264)
1234         return type == H264_NAL_SPS || type == H264_NAL_PPS;
1235     if (codec_id == AV_CODEC_ID_HEVC)
1236         return type == HEVC_NAL_VPS || type == HEVC_NAL_SPS || type == HEVC_NAL_PPS;
1237     return 0;
1238 }
1239
1240 static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
1241                                        CodedBitstreamFragment *frag)
1242 {
1243     uint8_t *data;
1244     size_t max_size, dp, sp;
1245     int err, i, zero_run;
1246
1247     for (i = 0; i < frag->nb_units; i++) {
1248         // Data should already all have been written when we get here.
1249         av_assert0(frag->units[i].data);
1250     }
1251
1252     max_size = 0;
1253     for (i = 0; i < frag->nb_units; i++) {
1254         // Start code + content with worst-case emulation prevention.
1255         max_size += 4 + frag->units[i].data_size * 3 / 2;
1256     }
1257
1258     data = av_realloc(NULL, max_size + AV_INPUT_BUFFER_PADDING_SIZE);
1259     if (!data)
1260         return AVERROR(ENOMEM);
1261
1262     dp = 0;
1263     for (i = 0; i < frag->nb_units; i++) {
1264         CodedBitstreamUnit *unit = &frag->units[i];
1265
1266         if (unit->data_bit_padding > 0) {
1267             if (i < frag->nb_units - 1)
1268                 av_log(ctx->log_ctx, AV_LOG_WARNING, "Probably invalid "
1269                        "unaligned padding on non-final NAL unit.\n");
1270             else
1271                 frag->data_bit_padding = unit->data_bit_padding;
1272         }
1273
1274         if (cbs_h2645_unit_requires_zero_byte(ctx->codec->codec_id, unit->type, i)) {
1275             // zero_byte
1276             data[dp++] = 0;
1277         }
1278         // start_code_prefix_one_3bytes
1279         data[dp++] = 0;
1280         data[dp++] = 0;
1281         data[dp++] = 1;
1282
1283         zero_run = 0;
1284         for (sp = 0; sp < unit->data_size; sp++) {
1285             if (zero_run < 2) {
1286                 if (unit->data[sp] == 0)
1287                     ++zero_run;
1288                 else
1289                     zero_run = 0;
1290             } else {
1291                 if ((unit->data[sp] & ~3) == 0) {
1292                     // emulation_prevention_three_byte
1293                     data[dp++] = 3;
1294                 }
1295                 zero_run = unit->data[sp] == 0;
1296             }
1297             data[dp++] = unit->data[sp];
1298         }
1299     }
1300
1301     av_assert0(dp <= max_size);
1302     err = av_reallocp(&data, dp + AV_INPUT_BUFFER_PADDING_SIZE);
1303     if (err)
1304         return err;
1305     memset(data + dp, 0, AV_INPUT_BUFFER_PADDING_SIZE);
1306
1307     frag->data_ref = av_buffer_create(data, dp + AV_INPUT_BUFFER_PADDING_SIZE,
1308                                       NULL, NULL, 0);
1309     if (!frag->data_ref) {
1310         av_freep(&data);
1311         return AVERROR(ENOMEM);
1312     }
1313
1314     frag->data = data;
1315     frag->data_size = dp;
1316
1317     return 0;
1318 }
1319
1320 static void cbs_h264_flush(CodedBitstreamContext *ctx)
1321 {
1322     CodedBitstreamH264Context *h264 = ctx->priv_data;
1323
1324     for (int i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++) {
1325         av_buffer_unref(&h264->sps_ref[i]);
1326         h264->sps[i] = NULL;
1327     }
1328     for (int i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++) {
1329         av_buffer_unref(&h264->pps_ref[i]);
1330         h264->pps[i] = NULL;
1331     }
1332
1333     h264->active_sps = NULL;
1334     h264->active_pps = NULL;
1335     h264->last_slice_nal_unit_type = 0;
1336 }
1337
1338 static void cbs_h264_close(CodedBitstreamContext *ctx)
1339 {
1340     CodedBitstreamH264Context *h264 = ctx->priv_data;
1341     int i;
1342
1343     ff_h2645_packet_uninit(&h264->common.read_packet);
1344
1345     for (i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++)
1346         av_buffer_unref(&h264->sps_ref[i]);
1347     for (i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++)
1348         av_buffer_unref(&h264->pps_ref[i]);
1349 }
1350
1351 static void cbs_h265_flush(CodedBitstreamContext *ctx)
1352 {
1353     CodedBitstreamH265Context *h265 = ctx->priv_data;
1354
1355     for (int i = 0; i < FF_ARRAY_ELEMS(h265->vps); i++) {
1356         av_buffer_unref(&h265->vps_ref[i]);
1357         h265->vps[i] = NULL;
1358     }
1359     for (int i = 0; i < FF_ARRAY_ELEMS(h265->sps); i++) {
1360         av_buffer_unref(&h265->sps_ref[i]);
1361         h265->sps[i] = NULL;
1362     }
1363     for (int i = 0; i < FF_ARRAY_ELEMS(h265->pps); i++) {
1364         av_buffer_unref(&h265->pps_ref[i]);
1365         h265->pps[i] = NULL;
1366     }
1367
1368     h265->active_vps = NULL;
1369     h265->active_sps = NULL;
1370     h265->active_pps = NULL;
1371 }
1372
1373 static void cbs_h265_close(CodedBitstreamContext *ctx)
1374 {
1375     CodedBitstreamH265Context *h265 = ctx->priv_data;
1376     int i;
1377
1378     ff_h2645_packet_uninit(&h265->common.read_packet);
1379
1380     for (i = 0; i < FF_ARRAY_ELEMS(h265->vps); i++)
1381         av_buffer_unref(&h265->vps_ref[i]);
1382     for (i = 0; i < FF_ARRAY_ELEMS(h265->sps); i++)
1383         av_buffer_unref(&h265->sps_ref[i]);
1384     for (i = 0; i < FF_ARRAY_ELEMS(h265->pps); i++)
1385         av_buffer_unref(&h265->pps_ref[i]);
1386 }
1387
1388 static void cbs_h264_free_sei(void *opaque, uint8_t *content)
1389 {
1390     H264RawSEI *sei = (H264RawSEI*)content;
1391     ff_cbs_sei_free_message_list(&sei->message_list);
1392     av_free(content);
1393 }
1394
1395 static const CodedBitstreamUnitTypeDescriptor cbs_h264_unit_types[] = {
1396     CBS_UNIT_TYPE_POD(H264_NAL_SPS,     H264RawSPS),
1397     CBS_UNIT_TYPE_POD(H264_NAL_SPS_EXT, H264RawSPSExtension),
1398
1399     CBS_UNIT_TYPE_INTERNAL_REF(H264_NAL_PPS, H264RawPPS, slice_group_id),
1400
1401     {
1402         .nb_unit_types  = 3,
1403         .unit_types     = {
1404             H264_NAL_IDR_SLICE,
1405             H264_NAL_SLICE,
1406             H264_NAL_AUXILIARY_SLICE,
1407         },
1408         .content_type   = CBS_CONTENT_TYPE_INTERNAL_REFS,
1409         .content_size   = sizeof(H264RawSlice),
1410         .nb_ref_offsets = 1,
1411         .ref_offsets    = { offsetof(H264RawSlice, data) },
1412     },
1413
1414     CBS_UNIT_TYPE_POD(H264_NAL_AUD,          H264RawAUD),
1415     CBS_UNIT_TYPE_POD(H264_NAL_FILLER_DATA,  H264RawFiller),
1416     CBS_UNIT_TYPE_POD(H264_NAL_END_SEQUENCE, H264RawNALUnitHeader),
1417     CBS_UNIT_TYPE_POD(H264_NAL_END_STREAM,   H264RawNALUnitHeader),
1418
1419     CBS_UNIT_TYPE_COMPLEX(H264_NAL_SEI, H264RawSEI, &cbs_h264_free_sei),
1420
1421     CBS_UNIT_TYPE_END_OF_LIST
1422 };
1423
1424 static void cbs_h265_free_sei(void *opaque, uint8_t *content)
1425 {
1426     H265RawSEI *sei = (H265RawSEI*)content;
1427     ff_cbs_sei_free_message_list(&sei->message_list);
1428     av_free(content);
1429 }
1430
1431 static const CodedBitstreamUnitTypeDescriptor cbs_h265_unit_types[] = {
1432     CBS_UNIT_TYPE_INTERNAL_REF(HEVC_NAL_VPS, H265RawVPS, extension_data.data),
1433     CBS_UNIT_TYPE_INTERNAL_REF(HEVC_NAL_SPS, H265RawSPS, extension_data.data),
1434     CBS_UNIT_TYPE_INTERNAL_REF(HEVC_NAL_PPS, H265RawPPS, extension_data.data),
1435
1436     CBS_UNIT_TYPE_POD(HEVC_NAL_AUD, H265RawAUD),
1437
1438     {
1439         // Slices of non-IRAP pictures.
1440         .nb_unit_types         = CBS_UNIT_TYPE_RANGE,
1441         .unit_type_range_start = HEVC_NAL_TRAIL_N,
1442         .unit_type_range_end   = HEVC_NAL_RASL_R,
1443
1444         .content_type   = CBS_CONTENT_TYPE_INTERNAL_REFS,
1445         .content_size   = sizeof(H265RawSlice),
1446         .nb_ref_offsets = 1,
1447         .ref_offsets    = { offsetof(H265RawSlice, data) },
1448     },
1449
1450     {
1451         // Slices of IRAP pictures.
1452         .nb_unit_types         = CBS_UNIT_TYPE_RANGE,
1453         .unit_type_range_start = HEVC_NAL_BLA_W_LP,
1454         .unit_type_range_end   = HEVC_NAL_CRA_NUT,
1455
1456         .content_type   = CBS_CONTENT_TYPE_INTERNAL_REFS,
1457         .content_size   = sizeof(H265RawSlice),
1458         .nb_ref_offsets = 1,
1459         .ref_offsets    = { offsetof(H265RawSlice, data) },
1460     },
1461
1462     {
1463         .nb_unit_types  = 2,
1464         .unit_types     = {
1465             HEVC_NAL_SEI_PREFIX,
1466             HEVC_NAL_SEI_SUFFIX
1467         },
1468         .content_type   = CBS_CONTENT_TYPE_COMPLEX,
1469         .content_size   = sizeof(H265RawSEI),
1470         .content_free   = &cbs_h265_free_sei,
1471     },
1472
1473     CBS_UNIT_TYPE_END_OF_LIST
1474 };
1475
1476 const CodedBitstreamType ff_cbs_type_h264 = {
1477     .codec_id          = AV_CODEC_ID_H264,
1478
1479     .priv_data_size    = sizeof(CodedBitstreamH264Context),
1480
1481     .unit_types        = cbs_h264_unit_types,
1482
1483     .split_fragment    = &cbs_h2645_split_fragment,
1484     .read_unit         = &cbs_h264_read_nal_unit,
1485     .write_unit        = &cbs_h264_write_nal_unit,
1486     .assemble_fragment = &cbs_h2645_assemble_fragment,
1487
1488     .flush             = &cbs_h264_flush,
1489     .close             = &cbs_h264_close,
1490 };
1491
1492 const CodedBitstreamType ff_cbs_type_h265 = {
1493     .codec_id          = AV_CODEC_ID_HEVC,
1494
1495     .priv_data_size    = sizeof(CodedBitstreamH265Context),
1496
1497     .unit_types        = cbs_h265_unit_types,
1498
1499     .split_fragment    = &cbs_h2645_split_fragment,
1500     .read_unit         = &cbs_h265_read_nal_unit,
1501     .write_unit        = &cbs_h265_write_nal_unit,
1502     .assemble_fragment = &cbs_h2645_assemble_fragment,
1503
1504     .flush             = &cbs_h265_flush,
1505     .close             = &cbs_h265_close,
1506 };
1507
1508 static const SEIMessageTypeDescriptor cbs_sei_common_types[] = {
1509     {
1510         SEI_TYPE_FILLER_PAYLOAD,
1511         1, 1,
1512         sizeof(SEIRawFillerPayload),
1513         SEI_MESSAGE_RW(sei, filler_payload),
1514     },
1515     {
1516         SEI_TYPE_USER_DATA_REGISTERED_ITU_T_T35,
1517         1, 1,
1518         sizeof(SEIRawUserDataRegistered),
1519         SEI_MESSAGE_RW(sei, user_data_registered),
1520     },
1521     {
1522         SEI_TYPE_USER_DATA_UNREGISTERED,
1523         1, 1,
1524         sizeof(SEIRawUserDataUnregistered),
1525         SEI_MESSAGE_RW(sei, user_data_unregistered),
1526     },
1527     {
1528         SEI_TYPE_MASTERING_DISPLAY_COLOUR_VOLUME,
1529         1, 0,
1530         sizeof(SEIRawMasteringDisplayColourVolume),
1531         SEI_MESSAGE_RW(sei, mastering_display_colour_volume),
1532     },
1533     {
1534         SEI_TYPE_CONTENT_LIGHT_LEVEL_INFO,
1535         1, 0,
1536         sizeof(SEIRawContentLightLevelInfo),
1537         SEI_MESSAGE_RW(sei, content_light_level_info),
1538     },
1539     {
1540         SEI_TYPE_ALTERNATIVE_TRANSFER_CHARACTERISTICS,
1541         1, 0,
1542         sizeof(SEIRawAlternativeTransferCharacteristics),
1543         SEI_MESSAGE_RW(sei, alternative_transfer_characteristics),
1544     },
1545     SEI_MESSAGE_TYPE_END,
1546 };
1547
1548 static const SEIMessageTypeDescriptor cbs_sei_h264_types[] = {
1549     {
1550         SEI_TYPE_BUFFERING_PERIOD,
1551         1, 0,
1552         sizeof(H264RawSEIBufferingPeriod),
1553         SEI_MESSAGE_RW(h264, sei_buffering_period),
1554     },
1555     {
1556         SEI_TYPE_PIC_TIMING,
1557         1, 0,
1558         sizeof(H264RawSEIPicTiming),
1559         SEI_MESSAGE_RW(h264, sei_pic_timing),
1560     },
1561     {
1562         SEI_TYPE_PAN_SCAN_RECT,
1563         1, 0,
1564         sizeof(H264RawSEIPanScanRect),
1565         SEI_MESSAGE_RW(h264, sei_pan_scan_rect),
1566     },
1567     {
1568         SEI_TYPE_RECOVERY_POINT,
1569         1, 0,
1570         sizeof(H264RawSEIRecoveryPoint),
1571         SEI_MESSAGE_RW(h264, sei_recovery_point),
1572     },
1573     {
1574         SEI_TYPE_DISPLAY_ORIENTATION,
1575         1, 0,
1576         sizeof(H264RawSEIDisplayOrientation),
1577         SEI_MESSAGE_RW(h264, sei_display_orientation),
1578     },
1579     SEI_MESSAGE_TYPE_END
1580 };
1581
1582 static const SEIMessageTypeDescriptor cbs_sei_h265_types[] = {
1583     {
1584         SEI_TYPE_BUFFERING_PERIOD,
1585         1, 0,
1586         sizeof(H265RawSEIBufferingPeriod),
1587         SEI_MESSAGE_RW(h265, sei_buffering_period),
1588     },
1589     {
1590         SEI_TYPE_PIC_TIMING,
1591         1, 0,
1592         sizeof(H265RawSEIPicTiming),
1593         SEI_MESSAGE_RW(h265, sei_pic_timing),
1594     },
1595     {
1596         SEI_TYPE_PAN_SCAN_RECT,
1597         1, 0,
1598         sizeof(H265RawSEIPanScanRect),
1599         SEI_MESSAGE_RW(h265, sei_pan_scan_rect),
1600     },
1601     {
1602         SEI_TYPE_RECOVERY_POINT,
1603         1, 0,
1604         sizeof(H265RawSEIRecoveryPoint),
1605         SEI_MESSAGE_RW(h265, sei_recovery_point),
1606     },
1607     {
1608         SEI_TYPE_DISPLAY_ORIENTATION,
1609         1, 0,
1610         sizeof(H265RawSEIDisplayOrientation),
1611         SEI_MESSAGE_RW(h265, sei_display_orientation),
1612     },
1613     {
1614         SEI_TYPE_ACTIVE_PARAMETER_SETS,
1615         1, 0,
1616         sizeof(H265RawSEIActiveParameterSets),
1617         SEI_MESSAGE_RW(h265, sei_active_parameter_sets),
1618     },
1619     {
1620         SEI_TYPE_DECODED_PICTURE_HASH,
1621         0, 1,
1622         sizeof(H265RawSEIDecodedPictureHash),
1623         SEI_MESSAGE_RW(h265, sei_decoded_picture_hash),
1624     },
1625     {
1626         SEI_TYPE_TIME_CODE,
1627         1, 0,
1628         sizeof(H265RawSEITimeCode),
1629         SEI_MESSAGE_RW(h265, sei_time_code),
1630     },
1631     {
1632         SEI_TYPE_ALPHA_CHANNEL_INFO,
1633         1, 0,
1634         sizeof(H265RawSEIAlphaChannelInfo),
1635         SEI_MESSAGE_RW(h265, sei_alpha_channel_info),
1636     },
1637     SEI_MESSAGE_TYPE_END
1638 };
1639
1640 const SEIMessageTypeDescriptor *ff_cbs_sei_find_type(CodedBitstreamContext *ctx,
1641                                                      int payload_type)
1642 {
1643     const SEIMessageTypeDescriptor *codec_list;
1644     int i;
1645
1646     for (i = 0; cbs_sei_common_types[i].type >= 0; i++) {
1647         if (cbs_sei_common_types[i].type == payload_type)
1648             return &cbs_sei_common_types[i];
1649     }
1650
1651     switch (ctx->codec->codec_id) {
1652     case AV_CODEC_ID_H264:
1653         codec_list = cbs_sei_h264_types;
1654         break;
1655     case AV_CODEC_ID_H265:
1656         codec_list = cbs_sei_h265_types;
1657         break;
1658     default:
1659         return NULL;
1660     }
1661
1662     for (i = 0; codec_list[i].type >= 0; i++) {
1663         if (codec_list[i].type == payload_type)
1664             return &codec_list[i];
1665     }
1666
1667     return NULL;
1668 }