15890f73d7d00b297ad1ba1133dd76f5fb5210d4
[ffmpeg.git] / libavcodec / hevc_mp4toannexb_bsf.c
1 /*
2  * HEVC MP4 to Annex B byte stream format filter
3  * copyright (c) 2015 Anton Khirnov
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/intreadwrite.h"
25 #include "libavutil/mem.h"
26
27 #include "avcodec.h"
28 #include "bytestream.h"
29 #include "hevc.h"
30
31 #define MIN_HEVCC_LENGTH 23
32
33 typedef struct HEVCBSFContext {
34     uint8_t  length_size;
35     int      extradata_parsed;
36
37     int logged_nonmp4_warning;
38 } HEVCBSFContext;
39
40 static int hevc_extradata_to_annexb(AVCodecContext *avctx)
41 {
42     GetByteContext gb;
43     int length_size, num_arrays, i, j;
44     int ret = 0;
45
46     uint8_t *new_extradata = NULL;
47     size_t   new_extradata_size = 0;;
48
49     bytestream2_init(&gb, avctx->extradata, avctx->extradata_size);
50
51     bytestream2_skip(&gb, 21);
52     length_size = (bytestream2_get_byte(&gb) & 3) + 1;
53     num_arrays  = bytestream2_get_byte(&gb);
54
55     for (i = 0; i < num_arrays; i++) {
56         int type = bytestream2_get_byte(&gb) & 0x3f;
57         int cnt  = bytestream2_get_be16(&gb);
58
59         if (!(type == NAL_VPS || type == NAL_SPS || type == NAL_PPS ||
60               type == NAL_SEI_PREFIX || type == NAL_SEI_SUFFIX)) {
61             av_log(avctx, AV_LOG_ERROR, "Invalid NAL unit type in extradata: %d\n",
62                    type);
63             ret = AVERROR_INVALIDDATA;
64             goto fail;
65         }
66
67         for (j = 0; j < cnt; j++) {
68             int nalu_len = bytestream2_get_be16(&gb);
69
70             if (4 + AV_INPUT_BUFFER_PADDING_SIZE + nalu_len > SIZE_MAX - new_extradata_size) {
71                 ret = AVERROR_INVALIDDATA;
72                 goto fail;
73             }
74             ret = av_reallocp(&new_extradata, new_extradata_size + nalu_len + 4 + AV_INPUT_BUFFER_PADDING_SIZE);
75             if (ret < 0)
76                 goto fail;
77
78             AV_WB32(new_extradata + new_extradata_size, 1); // add the startcode
79             bytestream2_get_buffer(&gb, new_extradata + new_extradata_size + 4, nalu_len);
80             new_extradata_size += 4 + nalu_len;
81             memset(new_extradata + new_extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
82         }
83     }
84
85     av_freep(&avctx->extradata);
86     avctx->extradata      = new_extradata;
87     avctx->extradata_size = new_extradata_size;
88
89     if (!new_extradata_size)
90         av_log(avctx, AV_LOG_WARNING, "No parameter sets in the extradata\n");
91
92     return length_size;
93 fail:
94     av_freep(&new_extradata);
95     return ret;
96 }
97
98 static int hevc_mp4toannexb_filter(AVBitStreamFilterContext *bsfc,
99                                    AVCodecContext *avctx, const char *args,
100                                    uint8_t **poutbuf, int *poutbuf_size,
101                                    const uint8_t *buf, int buf_size,
102                                    int keyframe)
103 {
104     HEVCBSFContext *ctx = bsfc->priv_data;
105     GetByteContext gb;
106
107     uint8_t *out = NULL;
108     size_t   out_size = 0;
109     int got_irap = 0;
110     int i, ret = 0;
111
112     if (!ctx->extradata_parsed) {
113         if (avctx->extradata_size < MIN_HEVCC_LENGTH ||
114             AV_RB24(avctx->extradata) == 1           ||
115             AV_RB32(avctx->extradata) == 1) {
116             if (!ctx->logged_nonmp4_warning) {
117                 av_log(avctx, AV_LOG_VERBOSE,
118                        "The input looks like it is Annex B already\n");
119                 ctx->logged_nonmp4_warning = 1;
120             }
121             *poutbuf      = buf;
122             *poutbuf_size = buf_size;
123             return 0;
124         }
125
126         ret = hevc_extradata_to_annexb(avctx);
127         if (ret < 0)
128             return ret;
129         ctx->length_size      = ret;
130         ctx->extradata_parsed = 1;
131     }
132
133     *poutbuf_size = 0;
134     *poutbuf      = NULL;
135
136     bytestream2_init(&gb, buf, buf_size);
137
138     while (bytestream2_get_bytes_left(&gb)) {
139         uint32_t nalu_size = 0;
140         int      nalu_type;
141         int is_irap, add_extradata, extra_size;
142
143         for (i = 0; i < ctx->length_size; i++)
144             nalu_size = (nalu_size << 8) | bytestream2_get_byte(&gb);
145
146         nalu_type = (bytestream2_peek_byte(&gb) >> 1) & 0x3f;
147
148         /* prepend extradata to IRAP frames */
149         is_irap       = nalu_type >= 16 && nalu_type <= 23;
150         add_extradata = is_irap && !got_irap;
151         extra_size    = add_extradata * avctx->extradata_size;
152         got_irap     |= is_irap;
153
154         if (SIZE_MAX - out_size < 4             ||
155             SIZE_MAX - out_size - 4 < nalu_size ||
156             SIZE_MAX - out_size - 4 - nalu_size < extra_size) {
157             ret = AVERROR_INVALIDDATA;
158             goto fail;
159         }
160
161         ret = av_reallocp(&out, out_size + 4 + nalu_size + extra_size);
162         if (ret < 0)
163             goto fail;
164
165         if (add_extradata)
166             memcpy(out + out_size, avctx->extradata, extra_size);
167         AV_WB32(out + out_size + extra_size, 1);
168         bytestream2_get_buffer(&gb, out + out_size + 4 + extra_size, nalu_size);
169         out_size += 4 + nalu_size + extra_size;
170     }
171
172     *poutbuf      = out;
173     *poutbuf_size = out_size;
174
175     return 1;
176
177 fail:
178     av_freep(&out);
179     return ret;
180 }
181
182 AVBitStreamFilter ff_hevc_mp4toannexb_bsf = {
183     "hevc_mp4toannexb",
184     sizeof(HEVCBSFContext),
185     hevc_mp4toannexb_filter,
186 };