path: root/libavcodec/hevc_mp4toannexb_bsf.c (plain)
blob: 09bce5b34c24d5fedd9ed12bedc7b216e1f9a319

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 "bsf.h"
29	#include "bytestream.h"
30	#include "hevc.h"
31
32	#define MIN_HEVCC_LENGTH 23
33
34	typedef struct HEVCBSFContext {
35	uint8_t length_size;
36	int extradata_parsed;
37	} HEVCBSFContext;
38
39	static int hevc_extradata_to_annexb(AVBSFContext *ctx)
40	{
41	GetByteContext gb;
42	int length_size, num_arrays, i, j;
43	int ret = 0;
44
45	uint8_t *new_extradata = NULL;
46	size_t new_extradata_size = 0;
47
48	bytestream2_init(&gb, ctx->par_in->extradata, ctx->par_in->extradata_size);
49
50	bytestream2_skip(&gb, 21);
51	length_size = (bytestream2_get_byte(&gb) & 3) + 1;
52	num_arrays = bytestream2_get_byte(&gb);
53
54	for (i = 0; i < num_arrays; i++) {
55	int type = bytestream2_get_byte(&gb) & 0x3f;
56	int cnt = bytestream2_get_be16(&gb);
57
58	if (!(type == HEVC_NAL_VPS || type == HEVC_NAL_SPS || type == HEVC_NAL_PPS ||
59	type == HEVC_NAL_SEI_PREFIX || type == HEVC_NAL_SEI_SUFFIX)) {
60	av_log(ctx, AV_LOG_ERROR, "Invalid NAL unit type in extradata: %d\n",
61	type);
62	ret = AVERROR_INVALIDDATA;
63	goto fail;
64	}
65
66	for (j = 0; j < cnt; j++) {
67	int nalu_len = bytestream2_get_be16(&gb);
68
69	if (4 + AV_INPUT_BUFFER_PADDING_SIZE + nalu_len > SIZE_MAX - new_extradata_size) {
70	ret = AVERROR_INVALIDDATA;
71	goto fail;
72	}
73	ret = av_reallocp(&new_extradata, new_extradata_size + nalu_len + 4 + AV_INPUT_BUFFER_PADDING_SIZE);
74	if (ret < 0)
75	goto fail;
76
77	AV_WB32(new_extradata + new_extradata_size, 1); // add the startcode
78	bytestream2_get_buffer(&gb, new_extradata + new_extradata_size + 4, nalu_len);
79	new_extradata_size += 4 + nalu_len;
80	memset(new_extradata + new_extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
81	}
82	}
83
84	av_freep(&ctx->par_out->extradata);
85	ctx->par_out->extradata = new_extradata;
86	ctx->par_out->extradata_size = new_extradata_size;
87
88	if (!new_extradata_size)
89	av_log(ctx, AV_LOG_WARNING, "No parameter sets in the extradata\n");
90
91	return length_size;
92	fail:
93	av_freep(&new_extradata);
94	return ret;
95	}
96
97	static int hevc_mp4toannexb_init(AVBSFContext *ctx)
98	{
99	HEVCBSFContext *s = ctx->priv_data;
100	int ret;
101
102	if (ctx->par_in->extradata_size < MIN_HEVCC_LENGTH ||
103	AV_RB24(ctx->par_in->extradata) == 1 ||
104	AV_RB32(ctx->par_in->extradata) == 1) {
105	av_log(ctx, AV_LOG_VERBOSE,
106	"The input looks like it is Annex B already\n");
107	} else {
108	ret = hevc_extradata_to_annexb(ctx);
109	if (ret < 0)
110	return ret;
111	s->length_size = ret;
112	s->extradata_parsed = 1;
113	}
114
115	return 0;
116	}
117
118	static int hevc_mp4toannexb_filter(AVBSFContext *ctx, AVPacket *out)
119	{
120	HEVCBSFContext *s = ctx->priv_data;
121	AVPacket *in;
122	GetByteContext gb;
123
124	int got_irap = 0;
125	int i, ret = 0;
126
127	ret = ff_bsf_get_packet(ctx, &in);
128	if (ret < 0)
129	return ret;
130
131	if (!s->extradata_parsed) {
132	av_packet_move_ref(out, in);
133	av_packet_free(&in);
134	return 0;
135	}
136
137	bytestream2_init(&gb, in->data, in->size);
138
139	while (bytestream2_get_bytes_left(&gb)) {
140	uint32_t nalu_size = 0;
141	int nalu_type;
142	int is_irap, add_extradata, extra_size, prev_size;
143
144	for (i = 0; i < s->length_size; i++)
145	nalu_size = (nalu_size << 8) | bytestream2_get_byte(&gb);
146
147	nalu_type = (bytestream2_peek_byte(&gb) >> 1) & 0x3f;
148
149	/* prepend extradata to IRAP frames */
150	is_irap = nalu_type >= 16 && nalu_type <= 23;
151	add_extradata = is_irap && !got_irap;
152	extra_size = add_extradata * ctx->par_out->extradata_size;
153	got_irap |= is_irap;
154
155	if (SIZE_MAX - nalu_size < 4 ||
156	SIZE_MAX - 4 - nalu_size < extra_size) {
157	ret = AVERROR_INVALIDDATA;
158	goto fail;
159	}
160
161	prev_size = out->size;
162
163	ret = av_grow_packet(out, 4 + nalu_size + extra_size);
164	if (ret < 0)
165	goto fail;
166
167	if (add_extradata)
168	memcpy(out->data + prev_size, ctx->par_out->extradata, extra_size);
169	AV_WB32(out->data + prev_size + extra_size, 1);
170	bytestream2_get_buffer(&gb, out->data + prev_size + 4 + extra_size, nalu_size);
171	}
172
173	ret = av_packet_copy_props(out, in);
174	if (ret < 0)
175	goto fail;
176
177	fail:
178	if (ret < 0)
179	av_packet_unref(out);
180	av_packet_free(&in);
181
182	return ret;
183	}
184
185	static const enum AVCodecID codec_ids[] = {
186	AV_CODEC_ID_HEVC, AV_CODEC_ID_NONE,
187	};
188
189	const AVBitStreamFilter ff_hevc_mp4toannexb_bsf = {
190	.name = "hevc_mp4toannexb",
191	.priv_data_size = sizeof(HEVCBSFContext),
192	.init = hevc_mp4toannexb_init,
193	.filter = hevc_mp4toannexb_filter,
194	.codec_ids = codec_ids,
195	};
196