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