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