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