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

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