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