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