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

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