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 "golomb.h"
  27 #include "h264.h"
  28 #include "h264_sei.h"
  29 #include "h2645_parse.h"
  30
  31
  32 static int cbs_read_ue_golomb(CodedBitstreamContext *ctx, BitstreamContext *bc,
  33                               const char *name, uint32_t *write_to,
  34                               uint32_t range_min, uint32_t range_max)
  35 {
  36     uint32_t value;
  37     int position;
  38
  39     if (ctx->trace_enable) {
  40         char bits[65];
  41         unsigned int k;
  42         int i, j;
  43
  44         position = bitstream_tell(bc);
  45
  46         for (i = 0; i < 32; i++) {
  47             k = bitstream_read_bit(bc);
  48             bits[i] = k ? '1' : '0';
  49             if (k)
  50                 break;
  51         }
  52         if (i >= 32) {
  53             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid ue-golomb "
  54                    "code found while reading %s: "
  55                    "more than 31 zeroes.\n", name);
  56             return AVERROR_INVALIDDATA;
  57         }
  58         value = 1;
  59         for (j = 0; j < i; j++) {
  60             k = bitstream_read_bit(bc);
  61             bits[i + j + 1] = k ? '1' : '0';
  62             value = value << 1 | k;
  63         }
  64         bits[i + j + 1] = 0;
  65         --value;
  66
  67         ff_cbs_trace_syntax_element(ctx, position, name, bits, value);
  68     } else {
  69         value = get_ue_golomb_long(bc);
  70     }
  71
  72     if (value < range_min || value > range_max) {
  73         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
  74                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
  75                name, value, range_min, range_max);
  76         return AVERROR_INVALIDDATA;
  77     }
  78
  79     *write_to = value;
  80     return 0;
  81 }
  82
  83 static int cbs_read_se_golomb(CodedBitstreamContext *ctx, BitstreamContext *bc,
  84                               const char *name, int32_t *write_to,
  85                               int32_t range_min, int32_t range_max)
  86 {
  87     int32_t value;
  88     int position;
  89
  90     if (ctx->trace_enable) {
  91         char bits[65];
  92         uint32_t v;
  93         unsigned int k;
  94         int i, j;
  95
  96         position = bitstream_tell(bc);
  97
  98         for (i = 0; i < 32; i++) {
  99             k = bitstream_read_bit(bc);
 100             bits[i] = k ? '1' : '0';
 101             if (k)
 102                 break;
 103         }
 104         if (i >= 32) {
 105             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid se-golomb "
 106                    "code found while reading %s: "
 107                    "more than 31 zeroes.\n", name);
 108             return AVERROR_INVALIDDATA;
 109         }
 110         v = 1;
 111         for (j = 0; j < i; j++) {
 112             k = bitstream_read_bit(bc);
 113             bits[i + j + 1] = k ? '1' : '0';
 114             v = v << 1 | k;
 115         }
 116         bits[i + j + 1] = 0;
 117         if (v & 1)
 118             value = -(int32_t)(v / 2);
 119         else
 120             value = v / 2;
 121
 122         ff_cbs_trace_syntax_element(ctx, position, name, bits, value);
 123     } else {
 124         value = get_se_golomb_long(bc);
 125     }
 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
 241
 242 #define READ
 243 #define READWRITE read
 244 #define RWContext BitstreamContext
 245
 246 #define xu(width, name, var, range_min, range_max) do { \
 247         uint32_t value = range_min; \
 248         CHECK(ff_cbs_read_unsigned(ctx, rw, width, #name, \
 249                                    &value, range_min, range_max)); \
 250         var = value; \
 251     } while (0)
 252 #define xue(name, var, range_min, range_max) do { \
 253         uint32_t value = range_min; \
 254         CHECK(cbs_read_ue_golomb(ctx, rw, #name, \
 255                                  &value, range_min, range_max)); \
 256         var = value; \
 257     } while (0)
 258 #define xse(name, var, range_min, range_max) do { \
 259         int32_t value = range_min; \
 260         CHECK(cbs_read_se_golomb(ctx, rw, #name, \
 261                                  &value, range_min, range_max)); \
 262         var = value; \
 263     } while (0)
 264
 265
 266 #define u(width, name, range_min, range_max) \
 267         xu(width, name, current->name, range_min, range_max)
 268 #define flag(name) u(1, name, 0, 1)
 269 #define ue(name, range_min, range_max) \
 270         xue(name, current->name, range_min, range_max)
 271 #define se(name, range_min, range_max) \
 272         xse(name, current->name, range_min, range_max)
 273
 274 #define infer(name, value) do { \
 275         current->name = value; \
 276     } while (0)
 277
 278 static int cbs_h2645_read_more_rbsp_data(BitstreamContext *bc)
 279 {
 280     int bits_left = bitstream_bits_left(bc);
 281     if (bits_left > 8)
 282         return 1;
 283     if (bitstream_peek(bc, bits_left) == 1 << (bits_left - 1))
 284         return 0;
 285     return 1;
 286 }
 287
 288 #define more_rbsp_data(var) ((var) = cbs_h2645_read_more_rbsp_data(rw))
 289
 290 #define byte_alignment(rw) (bitstream_tell(rw) % 8)
 291
 292 #define allocate(name, size) do { \
 293         name = av_mallocz(size); \
 294         if (!name) \
 295             return AVERROR(ENOMEM); \
 296     } while (0)
 297
 298 #define FUNC(name) FUNC_H264(READWRITE, name)
 299 #include "cbs_h264_syntax_template.c"
 300 #undef FUNC
 301
 302 #undef READ
 303 #undef READWRITE
 304 #undef RWContext
 305 #undef xu
 306 #undef xue
 307 #undef xse
 308 #undef u
 309 #undef flag
 310 #undef ue
 311 #undef se
 312 #undef infer
 313 #undef more_rbsp_data
 314 #undef byte_alignment
 315 #undef allocate
 316
 317
 318 #define WRITE
 319 #define READWRITE write
 320 #define RWContext PutBitContext
 321
 322 #define xu(width, name, var, range_min, range_max) do { \
 323         uint32_t value = var; \
 324         CHECK(ff_cbs_write_unsigned(ctx, rw, width, #name, \
 325                                     value, range_min, range_max)); \
 326     } while (0)
 327 #define xue(name, var, range_min, range_max) do { \
 328         uint32_t value = var; \
 329         CHECK(cbs_write_ue_golomb(ctx, rw, #name, \
 330                                   value, range_min, range_max)); \
 331     } while (0)
 332 #define xse(name, var, range_min, range_max) do { \
 333         int32_t value = var; \
 334         CHECK(cbs_write_se_golomb(ctx, rw, #name, \
 335                                   value, range_min, range_max)); \
 336     } while (0)
 337
 338 #define u(width, name, range_min, range_max) \
 339         xu(width, name, current->name, range_min, range_max)
 340 #define flag(name) u(1, name, 0, 1)
 341 #define ue(name, range_min, range_max) \
 342         xue(name, current->name, range_min, range_max)
 343 #define se(name, range_min, range_max) \
 344         xse(name, current->name, range_min, range_max)
 345
 346 #define infer(name, value) do { \
 347         if (current->name != (value)) { \
 348             av_log(ctx->log_ctx, AV_LOG_WARNING, "Warning: " \
 349                    "%s does not match inferred value: " \
 350                    "%"PRId64", but should be %"PRId64".\n", \
 351                    #name, (int64_t)current->name, (int64_t)(value)); \
 352         } \
 353     } while (0)
 354
 355 #define more_rbsp_data(var) (var)
 356
 357 #define byte_alignment(rw) (put_bits_count(rw) % 8)
 358
 359 #define allocate(name, size) do { \
 360         if (!name) { \
 361             av_log(ctx->log_ctx, AV_LOG_ERROR, "%s must be set " \
 362                    "for writing.\n", #name); \
 363             return AVERROR_INVALIDDATA; \
 364         } \
 365     } while (0)
 366
 367 #define FUNC(name) FUNC_H264(READWRITE, name)
 368 #include "cbs_h264_syntax_template.c"
 369 #undef FUNC
 370
 371 #undef WRITE
 372 #undef READWRITE
 373 #undef RWContext
 374 #undef xu
 375 #undef xue
 376 #undef xse
 377 #undef u
 378 #undef flag
 379 #undef ue
 380 #undef se
 381 #undef infer
 382 #undef more_rbsp_data
 383 #undef byte_alignment
 384 #undef allocate
 385
 386
 387 static void cbs_h264_free_sei(H264RawSEI *sei)
 388 {
 389     int i;
 390     for (i = 0; i < sei->payload_count; i++) {
 391         H264RawSEIPayload *payload = &sei->payload[i];
 392
 393         switch (payload->payload_type) {
 394         case H264_SEI_TYPE_BUFFERING_PERIOD:
 395         case H264_SEI_TYPE_PIC_TIMING:
 396         case H264_SEI_TYPE_RECOVERY_POINT:
 397         case H264_SEI_TYPE_DISPLAY_ORIENTATION:
 398             break;
 399         case H264_SEI_TYPE_USER_DATA_REGISTERED:
 400             av_freep(&payload->payload.user_data_registered.data);
 401             break;
 402         case H264_SEI_TYPE_USER_DATA_UNREGISTERED:
 403             av_freep(&payload->payload.user_data_unregistered.data);
 404             break;
 405         default:
 406             av_freep(&payload->payload.other.data);
 407             break;
 408         }
 409     }
 410 }
 411
 412 static void cbs_h264_free_slice(H264RawSlice *slice)
 413 {
 414     av_freep(&slice->data);
 415 }
 416
 417 static void cbs_h264_free_nal_unit(CodedBitstreamUnit *unit)
 418 {
 419     switch (unit->type) {
 420     case H264_NAL_SEI:
 421         cbs_h264_free_sei(unit->content);
 422         break;
 423     case H264_NAL_IDR_SLICE:
 424     case H264_NAL_SLICE:
 425         cbs_h264_free_slice(unit->content);
 426         break;
 427     }
 428     av_freep(&unit->content);
 429 }
 430
 431 static int cbs_h2645_fragment_add_nals(CodedBitstreamContext *ctx,
 432                                        CodedBitstreamFragment *frag,
 433                                        const H2645Packet *packet)
 434 {
 435     int err, i;
 436
 437     for (i = 0; i < packet->nb_nals; i++) {
 438         const H2645NAL *nal = &packet->nals[i];
 439         uint8_t *data;
 440
 441         data = av_malloc(nal->size);
 442         if (!data)
 443             return AVERROR(ENOMEM);
 444         memcpy(data, nal->data, nal->size);
 445
 446         err = ff_cbs_insert_unit_data(ctx, frag, -1, nal->type,
 447                                       data, nal->size);
 448         if (err < 0) {
 449             av_freep(&data);
 450             return err;
 451         }
 452     }
 453
 454     return 0;
 455 }
 456
 457 static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
 458                                     CodedBitstreamFragment *frag,
 459                                     int header)
 460 {
 461     enum AVCodecID codec_id = ctx->codec->codec_id;
 462     CodedBitstreamH2645Context *priv = ctx->priv_data;
 463     GetByteContext gbc;
 464     int err;
 465
 466     av_assert0(frag->data && frag->nb_units == 0);
 467     if (frag->data_size == 0)
 468         return 0;
 469
 470     if (header && frag->data[0] && codec_id == AV_CODEC_ID_H264) {
 471         // AVCC header.
 472         size_t size, start, end;
 473         int i, count, version;
 474
 475         priv->mp4 = 1;
 476
 477         bytestream2_init(&gbc, frag->data, frag->data_size);
 478
 479         if (bytestream2_get_bytes_left(&gbc) < 6)
 480             return AVERROR_INVALIDDATA;
 481
 482         version = bytestream2_get_byte(&gbc);
 483         if (version != 1) {
 484             av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid AVCC header: "
 485                    "first byte %u.", version);
 486             return AVERROR_INVALIDDATA;
 487         }
 488
 489         bytestream2_skip(&gbc, 3);
 490         priv->nal_length_size = (bytestream2_get_byte(&gbc) & 3) + 1;
 491
 492         // SPS array.
 493         count = bytestream2_get_byte(&gbc) & 0x1f;
 494         start = bytestream2_tell(&gbc);
 495         for (i = 0; i < count; i++) {
 496             if (bytestream2_get_bytes_left(&gbc) < 2 * (count - i))
 497                 return AVERROR_INVALIDDATA;
 498             size = bytestream2_get_be16(&gbc);
 499             if (bytestream2_get_bytes_left(&gbc) < size)
 500                 return AVERROR_INVALIDDATA;
 501             bytestream2_skip(&gbc, size);
 502         }
 503         end = bytestream2_tell(&gbc);
 504
 505         err = ff_h2645_packet_split(&priv->read_packet,
 506                                     frag->data + start, end - start,
 507                                     ctx->log_ctx, 1, 2, AV_CODEC_ID_H264);
 508         if (err < 0) {
 509             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC SPS array.\n");
 510             return err;
 511         }
 512         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 513         if (err < 0)
 514             return err;
 515
 516         // PPS array.
 517         count = bytestream2_get_byte(&gbc);
 518         start = bytestream2_tell(&gbc);
 519         for (i = 0; i < count; i++) {
 520             if (bytestream2_get_bytes_left(&gbc) < 2 * (count - i))
 521                 return AVERROR_INVALIDDATA;
 522             size = bytestream2_get_be16(&gbc);
 523             if (bytestream2_get_bytes_left(&gbc) < size)
 524                 return AVERROR_INVALIDDATA;
 525             bytestream2_skip(&gbc, size);
 526         }
 527         end = bytestream2_tell(&gbc);
 528
 529         err = ff_h2645_packet_split(&priv->read_packet,
 530                                     frag->data + start, end - start,
 531                                     ctx->log_ctx, 1, 2, AV_CODEC_ID_H264);
 532         if (err < 0) {
 533             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to split AVCC PPS array.\n");
 534             return err;
 535         }
 536         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 537         if (err < 0)
 538             return err;
 539
 540         if (bytestream2_get_bytes_left(&gbc) > 0) {
 541             av_log(ctx->log_ctx, AV_LOG_WARNING, "%u bytes left at end of AVCC "
 542                    "header.\n", bytestream2_get_bytes_left(&gbc));
 543         }
 544
 545     } else {
 546         // Annex B, or later MP4 with already-known parameters.
 547
 548         err = ff_h2645_packet_split(&priv->read_packet,
 549                                     frag->data, frag->data_size,
 550                                     ctx->log_ctx,
 551                                     priv->mp4, priv->nal_length_size,
 552                                     codec_id);
 553         if (err < 0)
 554             return err;
 555
 556         err = cbs_h2645_fragment_add_nals(ctx, frag, &priv->read_packet);
 557         if (err < 0)
 558             return err;
 559     }
 560
 561     return 0;
 562 }
 563
 564 #define cbs_h2645_replace_ps(h26n, ps_name, ps_var, id_element) \
 565 static int cbs_h26 ## h26n ## _replace_ ## ps_var(CodedBitstreamContext *ctx, \
 566                                                   const H26 ## h26n ## Raw ## ps_name *ps_var)  \
 567 { \
 568     CodedBitstreamH26 ## h26n ## Context *priv = ctx->priv_data; \
 569     unsigned int id = ps_var->id_element; \
 570     if (id > FF_ARRAY_ELEMS(priv->ps_var)) { \
 571         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid " #ps_name \
 572                " id : %d.\n", id); \
 573         return AVERROR_INVALIDDATA; \
 574     } \
 575     av_freep(&priv->ps_var[id]); \
 576     priv->ps_var[id] = av_malloc(sizeof(*ps_var)); \
 577     if (!priv->ps_var[id]) \
 578         return AVERROR(ENOMEM); \
 579     memcpy(priv->ps_var[id], ps_var, sizeof(*ps_var)); \
 580     return 0; \
 581 }
 582
 583 cbs_h2645_replace_ps(4, SPS, sps, seq_parameter_set_id)
 584 cbs_h2645_replace_ps(4, PPS, pps, pic_parameter_set_id)
 585
 586 static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
 587                                   CodedBitstreamUnit *unit)
 588 {
 589     BitstreamContext bc;
 590     int err;
 591
 592     err = bitstream_init(&bc, unit->data, 8 * unit->data_size);
 593     if (err < 0)
 594         return err;
 595
 596     switch (unit->type) {
 597     case H264_NAL_SPS:
 598         {
 599             H264RawSPS *sps;
 600
 601             sps = av_mallocz(sizeof(*sps));
 602             if (!sps)
 603                 return AVERROR(ENOMEM);
 604             err = cbs_h264_read_sps(ctx, &bc, sps);
 605             if (err >= 0)
 606                 err = cbs_h264_replace_sps(ctx, sps);
 607             if (err < 0) {
 608                 av_free(sps);
 609                 return err;
 610             }
 611
 612             unit->content = sps;
 613         }
 614         break;
 615
 616     case H264_NAL_SPS_EXT:
 617         {
 618             H264RawSPSExtension *sps_ext;
 619
 620             sps_ext = av_mallocz(sizeof(*sps_ext));
 621             if (!sps_ext)
 622                 return AVERROR(ENOMEM);
 623             err = cbs_h264_read_sps_extension(ctx, &bc, sps_ext);
 624             if (err < 0) {
 625                 av_free(sps_ext);
 626                 return err;
 627             }
 628
 629             unit->content = sps_ext;
 630         }
 631         break;
 632
 633     case H264_NAL_PPS:
 634         {
 635             H264RawPPS *pps;
 636
 637             pps = av_mallocz(sizeof(*pps));
 638             if (!pps)
 639                 return AVERROR(ENOMEM);
 640             err = cbs_h264_read_pps(ctx, &bc, pps);
 641             if (err >= 0)
 642                 err = cbs_h264_replace_pps(ctx, pps);
 643             if (err < 0) {
 644                 av_free(pps);
 645                 return err;
 646             }
 647
 648             unit->content = pps;
 649         }
 650         break;
 651
 652     case H264_NAL_SLICE:
 653     case H264_NAL_IDR_SLICE:
 654     case H264_NAL_AUXILIARY_SLICE:
 655         {
 656             H264RawSlice *slice;
 657             int pos, len;
 658
 659             slice = av_mallocz(sizeof(*slice));
 660             if (!slice)
 661                 return AVERROR(ENOMEM);
 662             err = cbs_h264_read_slice_header(ctx, &bc, &slice->header);
 663             if (err < 0) {
 664                 av_free(slice);
 665                 return err;
 666             }
 667
 668             pos = bitstream_tell(&bc);
 669             len = unit->data_size;
 670             if (!unit->data[len - 1]) {
 671                 int z;
 672                 for (z = 0; z < len && !unit->data[len - z - 1]; z++);
 673                 av_log(ctx->log_ctx, AV_LOG_DEBUG, "Deleted %d trailing zeroes "
 674                        "from slice data.\n", z);
 675                 len -= z;
 676             }
 677
 678             slice->data_size = len - pos / 8;
 679             slice->data = av_malloc(slice->data_size);
 680             if (!slice->data) {
 681                 av_free(slice);
 682                 return AVERROR(ENOMEM);
 683             }
 684             memcpy(slice->data,
 685                    unit->data + pos / 8, slice->data_size);
 686             slice->data_bit_start = pos % 8;
 687
 688             unit->content = slice;
 689         }
 690         break;
 691
 692     case H264_NAL_AUD:
 693         {
 694             H264RawAUD *aud;
 695
 696             aud = av_mallocz(sizeof(*aud));
 697             if (!aud)
 698                 return AVERROR(ENOMEM);
 699             err = cbs_h264_read_aud(ctx, &bc, aud);
 700             if (err < 0) {
 701                 av_free(aud);
 702                 return err;
 703             }
 704
 705             unit->content = aud;
 706         }
 707         break;
 708
 709     case H264_NAL_SEI:
 710         {
 711             H264RawSEI *sei;
 712
 713             sei = av_mallocz(sizeof(*sei));
 714             if (!sei)
 715                 return AVERROR(ENOMEM);
 716             err = cbs_h264_read_sei(ctx, &bc, sei);
 717             if (err < 0) {
 718                 cbs_h264_free_sei(sei);
 719                 return err;
 720             }
 721
 722             unit->content = sei;
 723         }
 724         break;
 725
 726     default:
 727         return AVERROR(ENOSYS);
 728     }
 729
 730     return 0;
 731 }
 732
 733 static int cbs_h264_write_nal_unit(CodedBitstreamContext *ctx,
 734                                    CodedBitstreamUnit *unit,
 735                                    PutBitContext *pbc)
 736 {
 737     int err;
 738
 739     switch (unit->type) {
 740     case H264_NAL_SPS:
 741         {
 742             H264RawSPS *sps = unit->content;
 743
 744             err = cbs_h264_write_sps(ctx, pbc, sps);
 745             if (err < 0)
 746                 return err;
 747
 748             err = cbs_h264_replace_sps(ctx, sps);
 749             if (err < 0)
 750                 return err;
 751         }
 752         break;
 753
 754     case H264_NAL_SPS_EXT:
 755         {
 756             H264RawSPSExtension *sps_ext;
 757
 758             err = cbs_h264_write_sps_extension(ctx, pbc, sps_ext);
 759             if (err < 0)
 760                 return err;
 761         }
 762         break;
 763
 764     case H264_NAL_PPS:
 765         {
 766             H264RawPPS *pps = unit->content;
 767
 768             err = cbs_h264_write_pps(ctx, pbc, pps);
 769             if (err < 0)
 770                 return err;
 771
 772             err = cbs_h264_replace_pps(ctx, pps);
 773             if (err < 0)
 774                 return err;
 775         }
 776         break;
 777
 778     case H264_NAL_SLICE:
 779     case H264_NAL_IDR_SLICE:
 780         {
 781             H264RawSlice *slice = unit->content;
 782             BitstreamContext bc;
 783             int bits_left, end, zeroes;
 784
 785             err = cbs_h264_write_slice_header(ctx, pbc, &slice->header);
 786             if (err < 0)
 787                 return err;
 788
 789             if (slice->data) {
 790                 if (slice->data_size * 8 + 8 > put_bits_left(pbc))
 791                     return AVERROR(ENOSPC);
 792
 793                 bitstream_init(&bc, slice->data, slice->data_size * 8);
 794                 bitstream_skip(&bc, slice->data_bit_start);
 795
 796                 // Copy in two-byte blocks, but stop before copying the
 797                 // rbsp_stop_one_bit in the final byte.
 798                 while (bitstream_bits_left(&bc) > 23)
 799                     put_bits(pbc, 16, bitstream_read(&bc, 16));
 800
 801                 bits_left = bitstream_bits_left(&bc);
 802                 end = bitstream_read(&bc, bits_left);
 803
 804                 // rbsp_stop_one_bit must be present here.
 805                 av_assert0(end);
 806                 zeroes = ff_ctz(end);
 807                 if (bits_left > zeroes + 1)
 808                     put_bits(pbc, bits_left - zeroes - 1,
 809                              end >> (zeroes + 1));
 810                 put_bits(pbc, 1, 1);
 811                 while (put_bits_count(pbc) % 8 != 0)
 812                     put_bits(pbc, 1, 0);
 813             } else {
 814                 // No slice data - that was just the header.
 815                 // (Bitstream may be unaligned!)
 816             }
 817         }
 818         break;
 819
 820     case H264_NAL_AUD:
 821         {
 822             err = cbs_h264_write_aud(ctx, pbc, unit->content);
 823             if (err < 0)
 824                 return err;
 825         }
 826         break;
 827
 828     case H264_NAL_SEI:
 829         {
 830             err = cbs_h264_write_sei(ctx, pbc, unit->content);
 831             if (err < 0)
 832                 return err;
 833         }
 834         break;
 835
 836     default:
 837         av_log(ctx->log_ctx, AV_LOG_ERROR, "Write unimplemented for "
 838                "NAL unit type %"PRIu32".\n", unit->type);
 839         return AVERROR_PATCHWELCOME;
 840     }
 841
 842     return 0;
 843 }
 844
 845 static int cbs_h2645_write_nal_unit(CodedBitstreamContext *ctx,
 846                                     CodedBitstreamUnit *unit)
 847 {
 848     CodedBitstreamH2645Context *priv = ctx->priv_data;
 849     enum AVCodecID codec_id = ctx->codec->codec_id;
 850     PutBitContext pbc;
 851     int err;
 852
 853     if (!priv->write_buffer) {
 854         // Initial write buffer size is 1MB.
 855         priv->write_buffer_size = 1024 * 1024;
 856
 857     reallocate_and_try_again:
 858         err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
 859         if (err < 0) {
 860             av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
 861                    "sufficiently large write buffer (last attempt "
 862                    "%zu bytes).\n", priv->write_buffer_size);
 863             return err;
 864         }
 865     }
 866
 867     init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
 868
 869     err = cbs_h264_write_nal_unit(ctx, unit, &pbc);
 870
 871     if (err == AVERROR(ENOSPC)) {
 872         // Overflow.
 873         priv->write_buffer_size *= 2;
 874         goto reallocate_and_try_again;
 875     }
 876     // Overflow but we didn't notice.
 877     av_assert0(put_bits_count(&pbc) <= 8 * priv->write_buffer_size);
 878
 879     if (put_bits_count(&pbc) % 8)
 880         unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
 881     else
 882         unit->data_bit_padding = 0;
 883
 884     unit->data_size = (put_bits_count(&pbc) + 7) / 8;
 885     flush_put_bits(&pbc);
 886
 887     err = av_reallocp(&unit->data, unit->data_size);
 888     if (err < 0)
 889         return err;
 890
 891     memcpy(unit->data, priv->write_buffer, unit->data_size);
 892
 893     return 0;
 894 }
 895
 896 static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
 897                                        CodedBitstreamFragment *frag)
 898 {
 899     uint8_t *data;
 900     size_t max_size, dp, sp;
 901     int err, i, zero_run;
 902
 903     for (i = 0; i < frag->nb_units; i++) {
 904         // Data should already all have been written when we get here.
 905         av_assert0(frag->units[i].data);
 906     }
 907
 908     max_size = 0;
 909     for (i = 0; i < frag->nb_units; i++) {
 910         // Start code + content with worst-case emulation prevention.
 911         max_size += 3 + frag->units[i].data_size * 3 / 2;
 912     }
 913
 914     data = av_malloc(max_size);
 915     if (!data)
 916         return AVERROR(ENOMEM);
 917
 918     dp = 0;
 919     for (i = 0; i < frag->nb_units; i++) {
 920         CodedBitstreamUnit *unit = &frag->units[i];
 921
 922         if (unit->data_bit_padding > 0) {
 923             if (i < frag->nb_units - 1)
 924                 av_log(ctx->log_ctx, AV_LOG_WARNING, "Probably invalid "
 925                        "unaligned padding on non-final NAL unit.\n");
 926             else
 927                 frag->data_bit_padding = unit->data_bit_padding;
 928         }
 929
 930         if (unit->type == H264_NAL_SPS ||
 931             unit->type == H264_NAL_PPS ||
 932             i == 0 /* (Assume this is the start of an access unit.) */) {
 933             // zero_byte
 934             data[dp++] = 0;
 935         }
 936         // start_code_prefix_one_3bytes
 937         data[dp++] = 0;
 938         data[dp++] = 0;
 939         data[dp++] = 1;
 940
 941         zero_run = 0;
 942         for (sp = 0; sp < unit->data_size; sp++) {
 943             if (zero_run < 2) {
 944                 if (unit->data[sp] == 0)
 945                     ++zero_run;
 946                 else
 947                     zero_run = 0;
 948             } else {
 949                 if ((unit->data[sp] & ~3) == 0) {
 950                     // emulation_prevention_three_byte
 951                     data[dp++] = 3;
 952                 }
 953                 zero_run = unit->data[sp] == 0;
 954             }
 955             data[dp++] = unit->data[sp];
 956         }
 957     }
 958
 959     av_assert0(dp <= max_size);
 960     err = av_reallocp(&data, dp);
 961     if (err)
 962         return err;
 963
 964     frag->data = data;
 965     frag->data_size = dp;
 966
 967     return 0;
 968 }
 969
 970 static void cbs_h264_close(CodedBitstreamContext *ctx)
 971 {
 972     CodedBitstreamH264Context *h264 = ctx->priv_data;
 973     int i;
 974
 975     ff_h2645_packet_uninit(&h264->common.read_packet);
 976
 977     av_freep(&h264->common.write_buffer);
 978
 979     for (i = 0; i < FF_ARRAY_ELEMS(h264->sps); i++)
 980         av_freep(&h264->sps[i]);
 981     for (i = 0; i < FF_ARRAY_ELEMS(h264->pps); i++)
 982         av_freep(&h264->pps[i]);
 983 }
 984
 985 const CodedBitstreamType ff_cbs_type_h264 = {
 986     .codec_id          = AV_CODEC_ID_H264,
 987
 988     .priv_data_size    = sizeof(CodedBitstreamH264Context),
 989
 990     .split_fragment    = &cbs_h2645_split_fragment,
 991     .read_unit         = &cbs_h264_read_nal_unit,
 992     .write_unit        = &cbs_h2645_write_nal_unit,
 993     .assemble_fragment = &cbs_h2645_assemble_fragment,
 994
 995     .free_unit         = &cbs_h264_free_nal_unit,
 996     .close             = &cbs_h264_close,
 997 };