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

   1 /*
   2  * H.264 MP4 to Annex B byte stream format filter
   3  * Copyright (c) 2007 Benoit Fouet <benoit.fouet@free.fr>
   4  *
   5  * This file is part of FFmpeg.
   6  *
   7  * FFmpeg is free software; you can redistribute it and/or
   8  * modify it under the terms of the GNU Lesser General Public
   9  * License as published by the Free Software Foundation; either
  10  * version 2.1 of the License, or (at your option) any later version.
  11  *
  12  * FFmpeg is distributed in the hope that it will be useful,
  13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15  * Lesser General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU Lesser General Public
  18  * License along with FFmpeg; if not, write to the Free Software
  19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  20  */
  21
  22 #include <string.h>
  23
  24 #include "libavutil/avassert.h"
  25 #include "libavutil/intreadwrite.h"
  26 #include "libavutil/mem.h"
  27
  28 #include "avcodec.h"
  29 #include "bsf.h"
  30 #include "bsf_internal.h"
  31 #include "bytestream.h"
  32 #include "h264.h"
  33
  34 typedef struct H264BSFContext {
  35     uint8_t *sps;
  36     uint8_t *pps;
  37     int      sps_size;
  38     int      pps_size;
  39     uint8_t  length_size;
  40     uint8_t  new_idr;
  41     uint8_t  idr_sps_seen;
  42     uint8_t  idr_pps_seen;
  43     int      extradata_parsed;
  44 } H264BSFContext;
  45
  46 static void count_or_copy(uint8_t **out, uint64_t *out_size,
  47                           const uint8_t *in, int in_size, int ps, int copy)
  48 {
  49     uint8_t start_code_size = ps < 0 ? 0 : *out_size == 0 || ps ? 4 : 3;
  50
  51     if (copy) {
  52         memcpy(*out + start_code_size, in, in_size);
  53         if (start_code_size == 4) {
  54             AV_WB32(*out, 1);
  55         } else if (start_code_size) {
  56             (*out)[0] =
  57             (*out)[1] = 0;
  58             (*out)[2] = 1;
  59         }
  60         *out  += start_code_size + in_size;
  61     }
  62     *out_size += start_code_size + in_size;
  63 }
  64
  65 static int h264_extradata_to_annexb(AVBSFContext *ctx, const int padding)
  66 {
  67     H264BSFContext *s = ctx->priv_data;
  68     GetByteContext ogb, *gb = &ogb;
  69     uint16_t unit_size;
  70     uint32_t total_size                 = 0;
  71     uint8_t *out                        = NULL, unit_nb, sps_done = 0;
  72     static const uint8_t nalu_header[4] = { 0, 0, 0, 1 };
  73     int length_size, pps_offset = 0;
  74
  75     bytestream2_init(gb, ctx->par_in->extradata, ctx->par_in->extradata_size);
  76
  77     bytestream2_skipu(gb, 4);
  78
  79     /* retrieve length coded size */
  80     length_size = (bytestream2_get_byteu(gb) & 0x3) + 1;
  81
  82     /* retrieve sps and pps unit(s) */
  83     unit_nb = bytestream2_get_byteu(gb) & 0x1f; /* number of sps unit(s) */
  84     if (!unit_nb) {
  85         goto pps;
  86     }
  87
  88     while (unit_nb--) {
  89         int err;
  90
  91         /* possible overread ok due to padding */
  92         unit_size   = bytestream2_get_be16u(gb);
  93         total_size += unit_size + 4;
  94         av_assert1(total_size <= INT_MAX - padding);
  95         if (bytestream2_get_bytes_left(gb) < unit_size + !sps_done) {
  96             av_log(ctx, AV_LOG_ERROR, "Global extradata truncated, "
  97                    "corrupted stream or invalid MP4/AVCC bitstream\n");
  98             av_free(out);
  99             return AVERROR_INVALIDDATA;
 100         }
 101         if ((err = av_reallocp(&out, total_size + padding)) < 0)
 102             return err;
 103         memcpy(out + total_size - unit_size - 4, nalu_header, 4);
 104         bytestream2_get_bufferu(gb, out + total_size - unit_size, unit_size);
 105 pps:
 106         if (!unit_nb && !sps_done++) {
 107             unit_nb = bytestream2_get_byteu(gb); /* number of pps unit(s) */
 108             pps_offset = total_size;
 109         }
 110     }
 111
 112     if (out)
 113         memset(out + total_size, 0, padding);
 114
 115     if (pps_offset) {
 116         s->sps      = out;
 117         s->sps_size = pps_offset;
 118     } else {
 119         av_log(ctx, AV_LOG_WARNING,
 120                "Warning: SPS NALU missing or invalid. "
 121                "The resulting stream may not play.\n");
 122     }
 123     if (pps_offset < total_size) {
 124         s->pps      = out + pps_offset;
 125         s->pps_size = total_size - pps_offset;
 126     } else {
 127         av_log(ctx, AV_LOG_WARNING,
 128                "Warning: PPS NALU missing or invalid. "
 129                "The resulting stream may not play.\n");
 130     }
 131
 132     av_freep(&ctx->par_out->extradata);
 133     ctx->par_out->extradata      = out;
 134     ctx->par_out->extradata_size = total_size;
 135
 136     return length_size;
 137 }
 138
 139 static int h264_mp4toannexb_init(AVBSFContext *ctx)
 140 {
 141     H264BSFContext *s = ctx->priv_data;
 142     int extra_size = ctx->par_in->extradata_size;
 143     int ret;
 144
 145     /* retrieve sps and pps NAL units from extradata */
 146     if (!extra_size                                               ||
 147         (extra_size >= 3 && AV_RB24(ctx->par_in->extradata) == 1) ||
 148         (extra_size >= 4 && AV_RB32(ctx->par_in->extradata) == 1)) {
 149         av_log(ctx, AV_LOG_VERBOSE,
 150                "The input looks like it is Annex B already\n");
 151     } else if (extra_size >= 7) {
 152         ret = h264_extradata_to_annexb(ctx, AV_INPUT_BUFFER_PADDING_SIZE);
 153         if (ret < 0)
 154             return ret;
 155
 156         s->length_size      = ret;
 157         s->new_idr          = 1;
 158         s->idr_sps_seen     = 0;
 159         s->idr_pps_seen     = 0;
 160         s->extradata_parsed = 1;
 161     } else {
 162         av_log(ctx, AV_LOG_ERROR, "Invalid extradata size: %d\n", extra_size);
 163         return AVERROR_INVALIDDATA;
 164     }
 165
 166     return 0;
 167 }
 168
 169 static int h264_mp4toannexb_filter(AVBSFContext *ctx, AVPacket *opkt)
 170 {
 171     H264BSFContext *s = ctx->priv_data;
 172     AVPacket *in;
 173     uint8_t unit_type, new_idr, sps_seen, pps_seen;
 174     const uint8_t *buf;
 175     const uint8_t *buf_end;
 176     uint8_t *out;
 177     uint64_t out_size;
 178     int ret;
 179
 180     ret = ff_bsf_get_packet(ctx, &in);
 181     if (ret < 0)
 182         return ret;
 183
 184     /* nothing to filter */
 185     if (!s->extradata_parsed) {
 186         av_packet_move_ref(opkt, in);
 187         av_packet_free(&in);
 188         return 0;
 189     }
 190
 191     buf_end  = in->data + in->size;
 192
 193 #define LOG_ONCE(...) \
 194     if (j) \
 195         av_log(__VA_ARGS__)
 196     for (int j = 0; j < 2; j++) {
 197         buf      = in->data;
 198         new_idr  = s->new_idr;
 199         sps_seen = s->idr_sps_seen;
 200         pps_seen = s->idr_pps_seen;
 201         out_size = 0;
 202
 203         do {
 204             uint32_t nal_size = 0;
 205
 206             /* possible overread ok due to padding */
 207             for (int i = 0; i < s->length_size; i++)
 208                 nal_size = (nal_size << 8) | buf[i];
 209
 210             buf += s->length_size;
 211
 212             /* This check requires the cast as the right side might
 213              * otherwise be promoted to an unsigned value. */
 214             if ((int64_t)nal_size > buf_end - buf) {
 215                 ret = AVERROR_INVALIDDATA;
 216                 goto fail;
 217             }
 218
 219             if (!nal_size)
 220                 continue;
 221
 222             unit_type = *buf & 0x1f;
 223
 224             if (unit_type == H264_NAL_SPS) {
 225                 sps_seen = new_idr = 1;
 226             } else if (unit_type == H264_NAL_PPS) {
 227                 pps_seen = new_idr = 1;
 228                 /* if SPS has not been seen yet, prepend the AVCC one to PPS */
 229                 if (!sps_seen) {
 230                     if (!s->sps_size) {
 231                         LOG_ONCE(ctx, AV_LOG_WARNING, "SPS not present in the stream, nor in AVCC, stream may be unreadable\n");
 232                     } else {
 233                         count_or_copy(&out, &out_size, s->sps, s->sps_size, -1, j);
 234                         sps_seen = 1;
 235                     }
 236                 }
 237             }
 238
 239             /* If this is a new IDR picture following an IDR picture, reset the idr flag.
 240              * Just check first_mb_in_slice to be 0 as this is the simplest solution.
 241              * This could be checking idr_pic_id instead, but would complexify the parsing. */
 242             if (!new_idr && unit_type == H264_NAL_IDR_SLICE && (buf[1] & 0x80))
 243                 new_idr = 1;
 244
 245             /* prepend only to the first type 5 NAL unit of an IDR picture, if no sps/pps are already present */
 246             if (new_idr && unit_type == H264_NAL_IDR_SLICE && !sps_seen && !pps_seen) {
 247                 if (ctx->par_out->extradata)
 248                     count_or_copy(&out, &out_size, ctx->par_out->extradata,
 249                                   ctx->par_out->extradata_size, -1, j);
 250                 new_idr = 0;
 251             /* if only SPS has been seen, also insert PPS */
 252             } else if (new_idr && unit_type == H264_NAL_IDR_SLICE && sps_seen && !pps_seen) {
 253                 if (!s->pps_size) {
 254                     LOG_ONCE(ctx, AV_LOG_WARNING, "PPS not present in the stream, nor in AVCC, stream may be unreadable\n");
 255                 } else {
 256                     count_or_copy(&out, &out_size, s->pps, s->pps_size, -1, j);
 257                 }
 258             }
 259
 260             count_or_copy(&out, &out_size, buf, nal_size,
 261                           unit_type == H264_NAL_SPS || unit_type == H264_NAL_PPS, j);
 262             if (!new_idr && unit_type == H264_NAL_SLICE) {
 263                 new_idr  = 1;
 264                 sps_seen = 0;
 265                 pps_seen = 0;
 266             }
 267
 268             buf += nal_size;
 269         } while (buf < buf_end);
 270
 271         if (!j) {
 272             if (out_size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE) {
 273                 ret = AVERROR_INVALIDDATA;
 274                 goto fail;
 275             }
 276             ret = av_new_packet(opkt, out_size);
 277             if (ret < 0)
 278                 goto fail;
 279             out = opkt->data;
 280         }
 281     }
 282 #undef LOG_ONCE
 283
 284     av_assert1(out_size == opkt->size);
 285
 286     s->new_idr      = new_idr;
 287     s->idr_sps_seen = sps_seen;
 288     s->idr_pps_seen = pps_seen;
 289
 290     ret = av_packet_copy_props(opkt, in);
 291     if (ret < 0)
 292         goto fail;
 293
 294 fail:
 295     if (ret < 0)
 296         av_packet_unref(opkt);
 297     av_packet_free(&in);
 298
 299     return ret;
 300 }
 301
 302 static void h264_mp4toannexb_flush(AVBSFContext *ctx)
 303 {
 304     H264BSFContext *s = ctx->priv_data;
 305
 306     s->idr_sps_seen = 0;
 307     s->idr_pps_seen = 0;
 308     s->new_idr      = s->extradata_parsed;
 309 }
 310
 311 static const enum AVCodecID codec_ids[] = {
 312     AV_CODEC_ID_H264, AV_CODEC_ID_NONE,
 313 };
 314
 315 const AVBitStreamFilter ff_h264_mp4toannexb_bsf = {
 316     .name           = "h264_mp4toannexb",
 317     .priv_data_size = sizeof(H264BSFContext),
 318     .init           = h264_mp4toannexb_init,
 319     .filter         = h264_mp4toannexb_filter,
 320     .flush          = h264_mp4toannexb_flush,
 321     .codec_ids      = codec_ids,
 322 };