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