]> git.sesse.net Git - ffmpeg/blob - libavcodec/extract_extradata_bsf.c
avcodec/ac3enc: Use specific error code in validate_float_option()
[ffmpeg] / libavcodec / extract_extradata_bsf.c
1 /*
2  * This file is part of FFmpeg.
3  *
4  * FFmpeg 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  * FFmpeg 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 FFmpeg; if not, write to the Free Software
16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
17  */
18
19 #include <stdint.h>
20
21 #include "libavutil/common.h"
22 #include "libavutil/intreadwrite.h"
23 #include "libavutil/log.h"
24 #include "libavutil/opt.h"
25
26 #include "avcodec.h"
27 #include "av1.h"
28 #include "av1_parse.h"
29 #include "bsf.h"
30 #include "h2645_parse.h"
31 #include "h264.h"
32 #include "hevc.h"
33 #include "vc1_common.h"
34
35 typedef struct ExtractExtradataContext {
36     const AVClass *class;
37
38     int (*extract)(AVBSFContext *ctx, AVPacket *pkt,
39                    uint8_t **data, int *size);
40
41     /* AV1 specifc fields */
42     AV1Packet av1_pkt;
43
44     /* H264/HEVC specifc fields */
45     H2645Packet h2645_pkt;
46
47     /* AVOptions */
48     int remove;
49 } ExtractExtradataContext;
50
51 static int val_in_array(const int *arr, int len, int val)
52 {
53     int i;
54     for (i = 0; i < len; i++)
55         if (arr[i] == val)
56             return 1;
57     return 0;
58 }
59
60 static int extract_extradata_av1(AVBSFContext *ctx, AVPacket *pkt,
61                                  uint8_t **data, int *size)
62 {
63     static const int extradata_obu_types[] = {
64         AV1_OBU_SEQUENCE_HEADER, AV1_OBU_METADATA,
65     };
66     ExtractExtradataContext *s = ctx->priv_data;
67
68     int extradata_size = 0, filtered_size = 0;
69     int nb_extradata_obu_types = FF_ARRAY_ELEMS(extradata_obu_types);
70     int i, has_seq = 0, ret = 0;
71
72     ret = ff_av1_packet_split(&s->av1_pkt, pkt->data, pkt->size, ctx);
73     if (ret < 0)
74         return ret;
75
76     for (i = 0; i < s->av1_pkt.nb_obus; i++) {
77         AV1OBU *obu = &s->av1_pkt.obus[i];
78         if (val_in_array(extradata_obu_types, nb_extradata_obu_types, obu->type)) {
79             extradata_size += obu->raw_size;
80             if (obu->type == AV1_OBU_SEQUENCE_HEADER)
81                 has_seq = 1;
82         } else if (s->remove) {
83             filtered_size += obu->raw_size;
84         }
85     }
86
87     if (extradata_size && has_seq) {
88         AVBufferRef *filtered_buf;
89         uint8_t *extradata, *filtered_data;
90
91         if (s->remove) {
92             filtered_buf = av_buffer_alloc(filtered_size + AV_INPUT_BUFFER_PADDING_SIZE);
93             if (!filtered_buf) {
94                 return AVERROR(ENOMEM);
95             }
96             memset(filtered_buf->data + filtered_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
97
98             filtered_data = filtered_buf->data;
99         }
100
101         extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
102         if (!extradata) {
103             av_buffer_unref(&filtered_buf);
104             return AVERROR(ENOMEM);
105         }
106         memset(extradata + extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
107
108         *data = extradata;
109         *size = extradata_size;
110
111         for (i = 0; i < s->av1_pkt.nb_obus; i++) {
112             AV1OBU *obu = &s->av1_pkt.obus[i];
113             if (val_in_array(extradata_obu_types, nb_extradata_obu_types,
114                              obu->type)) {
115                 memcpy(extradata, obu->raw_data, obu->raw_size);
116                 extradata += obu->raw_size;
117             } else if (s->remove) {
118                 memcpy(filtered_data, obu->raw_data, obu->raw_size);
119                 filtered_data += obu->raw_size;
120             }
121         }
122
123         if (s->remove) {
124             av_buffer_unref(&pkt->buf);
125             pkt->buf  = filtered_buf;
126             pkt->data = filtered_buf->data;
127             pkt->size = filtered_size;
128         }
129     }
130
131     return 0;
132 }
133
134 static int extract_extradata_h2645(AVBSFContext *ctx, AVPacket *pkt,
135                                    uint8_t **data, int *size)
136 {
137     static const int extradata_nal_types_hevc[] = {
138         HEVC_NAL_VPS, HEVC_NAL_SPS, HEVC_NAL_PPS,
139     };
140     static const int extradata_nal_types_h264[] = {
141         H264_NAL_SPS, H264_NAL_PPS,
142     };
143
144     ExtractExtradataContext *s = ctx->priv_data;
145
146     int extradata_size = 0, filtered_size = 0;
147     const int *extradata_nal_types;
148     int nb_extradata_nal_types;
149     int i, has_sps = 0, has_vps = 0, ret = 0;
150
151     if (ctx->par_in->codec_id == AV_CODEC_ID_HEVC) {
152         extradata_nal_types    = extradata_nal_types_hevc;
153         nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_hevc);
154     } else {
155         extradata_nal_types    = extradata_nal_types_h264;
156         nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_h264);
157     }
158
159     ret = ff_h2645_packet_split(&s->h2645_pkt, pkt->data, pkt->size,
160                                 ctx, 0, 0, ctx->par_in->codec_id, 1);
161     if (ret < 0)
162         return ret;
163
164     for (i = 0; i < s->h2645_pkt.nb_nals; i++) {
165         H2645NAL *nal = &s->h2645_pkt.nals[i];
166         if (val_in_array(extradata_nal_types, nb_extradata_nal_types, nal->type)) {
167             extradata_size += nal->raw_size + 3;
168             if (ctx->par_in->codec_id == AV_CODEC_ID_HEVC) {
169                 if (nal->type == HEVC_NAL_SPS) has_sps = 1;
170                 if (nal->type == HEVC_NAL_VPS) has_vps = 1;
171             } else {
172                 if (nal->type == H264_NAL_SPS) has_sps = 1;
173             }
174         } else if (s->remove) {
175             filtered_size += nal->raw_size + 3;
176         }
177     }
178
179     if (extradata_size &&
180         ((ctx->par_in->codec_id == AV_CODEC_ID_HEVC && has_sps && has_vps) ||
181          (ctx->par_in->codec_id == AV_CODEC_ID_H264 && has_sps))) {
182         AVBufferRef *filtered_buf;
183         uint8_t *extradata, *filtered_data;
184
185         if (s->remove) {
186             filtered_buf = av_buffer_alloc(filtered_size + AV_INPUT_BUFFER_PADDING_SIZE);
187             if (!filtered_buf) {
188                 return AVERROR(ENOMEM);
189             }
190             memset(filtered_buf->data + filtered_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
191
192             filtered_data = filtered_buf->data;
193         }
194
195         extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
196         if (!extradata) {
197             av_buffer_unref(&filtered_buf);
198             return AVERROR(ENOMEM);
199         }
200         memset(extradata + extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
201
202         *data = extradata;
203         *size = extradata_size;
204
205         for (i = 0; i < s->h2645_pkt.nb_nals; i++) {
206             H2645NAL *nal = &s->h2645_pkt.nals[i];
207             if (val_in_array(extradata_nal_types, nb_extradata_nal_types,
208                              nal->type)) {
209                 AV_WB24(extradata, 1); // startcode
210                 memcpy(extradata + 3, nal->raw_data, nal->raw_size);
211                 extradata += 3 + nal->raw_size;
212             } else if (s->remove) {
213                 AV_WB24(filtered_data, 1); // startcode
214                 memcpy(filtered_data + 3, nal->raw_data, nal->raw_size);
215                 filtered_data += 3 + nal->raw_size;
216             }
217         }
218
219         if (s->remove) {
220             av_buffer_unref(&pkt->buf);
221             pkt->buf  = filtered_buf;
222             pkt->data = filtered_buf->data;
223             pkt->size = filtered_size;
224         }
225     }
226
227     return 0;
228 }
229
230 static int extract_extradata_vc1(AVBSFContext *ctx, AVPacket *pkt,
231                                  uint8_t **data, int *size)
232 {
233     ExtractExtradataContext *s = ctx->priv_data;
234     const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size;
235     uint32_t state = UINT32_MAX;
236     int has_extradata = 0, extradata_size = 0;
237
238     while (ptr < end) {
239         ptr = avpriv_find_start_code(ptr, end, &state);
240         if (state == VC1_CODE_SEQHDR || state == VC1_CODE_ENTRYPOINT) {
241             has_extradata = 1;
242         } else if (has_extradata && IS_MARKER(state)) {
243             extradata_size = ptr - 4 - pkt->data;
244             break;
245         }
246     }
247
248     if (extradata_size) {
249         *data = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
250         if (!*data)
251             return AVERROR(ENOMEM);
252
253         memcpy(*data, pkt->data, extradata_size);
254         memset(*data + extradata_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
255         *size = extradata_size;
256
257         if (s->remove) {
258             pkt->data += extradata_size;
259             pkt->size -= extradata_size;
260         }
261     }
262
263     return 0;
264 }
265
266 static int extract_extradata_mpeg12(AVBSFContext *ctx, AVPacket *pkt,
267                                      uint8_t **data, int *size)
268 {
269     ExtractExtradataContext *s = ctx->priv_data;
270     uint32_t state = UINT32_MAX;
271     int i, found = 0;
272
273     for (i = 0; i < pkt->size; i++) {
274         state = (state << 8) | pkt->data[i];
275         if (state == 0x1B3)
276             found = 1;
277         else if (found && state != 0x1B5 && state < 0x200 && state >= 0x100) {
278             if (i > 3) {
279                 *size = i - 3;
280                 *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
281                 if (!*data)
282                     return AVERROR(ENOMEM);
283
284                 memcpy(*data, pkt->data, *size);
285                 memset(*data + *size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
286
287                 if (s->remove) {
288                     pkt->data += *size;
289                     pkt->size -= *size;
290                 }
291             }
292             break;
293         }
294     }
295     return 0;
296 }
297
298 static int extract_extradata_mpeg4(AVBSFContext *ctx, AVPacket *pkt,
299                                    uint8_t **data, int *size)
300 {
301     ExtractExtradataContext *s = ctx->priv_data;
302     const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size;
303     uint32_t state = UINT32_MAX;
304
305     while (ptr < end) {
306         ptr = avpriv_find_start_code(ptr, end, &state);
307         if (state == 0x1B3 || state == 0x1B6) {
308             if (ptr - pkt->data > 4) {
309                 *size = ptr - 4 - pkt->data;
310                 *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
311                 if (!*data)
312                     return AVERROR(ENOMEM);
313
314                 memcpy(*data, pkt->data, *size);
315                 memset(*data + *size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
316
317                 if (s->remove) {
318                     pkt->data += *size;
319                     pkt->size -= *size;
320                 }
321             }
322             break;
323         }
324     }
325     return 0;
326 }
327
328 static const struct {
329     enum AVCodecID id;
330     int (*extract)(AVBSFContext *ctx, AVPacket *pkt,
331                    uint8_t **data, int *size);
332 } extract_tab[] = {
333     { AV_CODEC_ID_AV1,        extract_extradata_av1     },
334     { AV_CODEC_ID_AVS2,       extract_extradata_mpeg4   },
335     { AV_CODEC_ID_CAVS,       extract_extradata_mpeg4   },
336     { AV_CODEC_ID_H264,       extract_extradata_h2645   },
337     { AV_CODEC_ID_HEVC,       extract_extradata_h2645   },
338     { AV_CODEC_ID_MPEG1VIDEO, extract_extradata_mpeg12  },
339     { AV_CODEC_ID_MPEG2VIDEO, extract_extradata_mpeg12  },
340     { AV_CODEC_ID_MPEG4,      extract_extradata_mpeg4   },
341     { AV_CODEC_ID_VC1,        extract_extradata_vc1     },
342 };
343
344 static int extract_extradata_init(AVBSFContext *ctx)
345 {
346     ExtractExtradataContext *s = ctx->priv_data;
347     int i;
348
349     for (i = 0; i < FF_ARRAY_ELEMS(extract_tab); i++) {
350         if (extract_tab[i].id == ctx->par_in->codec_id) {
351             s->extract = extract_tab[i].extract;
352             break;
353         }
354     }
355     if (!s->extract)
356         return AVERROR_BUG;
357
358     return 0;
359 }
360
361 static int extract_extradata_filter(AVBSFContext *ctx, AVPacket *pkt)
362 {
363     ExtractExtradataContext *s = ctx->priv_data;
364     uint8_t *extradata = NULL;
365     int extradata_size;
366     int ret = 0;
367
368     ret = ff_bsf_get_packet_ref(ctx, pkt);
369     if (ret < 0)
370         return ret;
371
372     ret = s->extract(ctx, pkt, &extradata, &extradata_size);
373     if (ret < 0)
374         goto fail;
375
376     if (extradata) {
377         ret = av_packet_add_side_data(pkt, AV_PKT_DATA_NEW_EXTRADATA,
378                                       extradata, extradata_size);
379         if (ret < 0) {
380             av_freep(&extradata);
381             goto fail;
382         }
383     }
384
385     return 0;
386
387 fail:
388     av_packet_unref(pkt);
389     return ret;
390 }
391
392 static void extract_extradata_close(AVBSFContext *ctx)
393 {
394     ExtractExtradataContext *s = ctx->priv_data;
395     ff_av1_packet_uninit(&s->av1_pkt);
396     ff_h2645_packet_uninit(&s->h2645_pkt);
397 }
398
399 static const enum AVCodecID codec_ids[] = {
400     AV_CODEC_ID_AV1,
401     AV_CODEC_ID_AVS2,
402     AV_CODEC_ID_CAVS,
403     AV_CODEC_ID_H264,
404     AV_CODEC_ID_HEVC,
405     AV_CODEC_ID_MPEG1VIDEO,
406     AV_CODEC_ID_MPEG2VIDEO,
407     AV_CODEC_ID_MPEG4,
408     AV_CODEC_ID_VC1,
409     AV_CODEC_ID_NONE,
410 };
411
412 #define OFFSET(x) offsetof(ExtractExtradataContext, x)
413 #define FLAGS (AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_BSF_PARAM)
414 static const AVOption options[] = {
415     { "remove", "remove the extradata from the bitstream", OFFSET(remove), AV_OPT_TYPE_INT,
416         { .i64 = 0 }, 0, 1, FLAGS },
417     { NULL },
418 };
419
420 static const AVClass extract_extradata_class = {
421     .class_name = "extract_extradata",
422     .item_name  = av_default_item_name,
423     .option     = options,
424     .version    = LIBAVUTIL_VERSION_INT,
425 };
426
427 const AVBitStreamFilter ff_extract_extradata_bsf = {
428     .name           = "extract_extradata",
429     .codec_ids      = codec_ids,
430     .priv_data_size = sizeof(ExtractExtradataContext),
431     .priv_class     = &extract_extradata_class,
432     .init           = extract_extradata_init,
433     .filter         = extract_extradata_filter,
434     .close          = extract_extradata_close,
435 };