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