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

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