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