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/log.h"
22#include "libavutil/opt.h"
23
24#include "av1.h"
25#include "av1_parse.h"
26#include "bsf.h"
27#include "bsf_internal.h"
28#include "bytestream.h"
29#include "h2645_parse.h"
30#include "h264.h"
31#include "hevc.h"
32#include "startcode.h"
33#include "vc1_common.h"
34
35typedef struct ExtractExtradataContext {
36    const AVClass *class;
37
38    int (*extract)(AVBSFContext *ctx, AVPacket *pkt,
39                   uint8_t **data, int *size);
40
41    /* AV1 specific fields */
42    AV1Packet av1_pkt;
43
44    /* H264/HEVC specific fields */
45    H2645Packet h2645_pkt;
46
47    /* AVOptions */
48    int remove;
49} ExtractExtradataContext;
50
51static 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
60static 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 = NULL;
89        PutByteContext pb_filtered_data, pb_extradata;
90        uint8_t *extradata;
91
92        if (s->remove) {
93            filtered_buf = av_buffer_alloc(filtered_size + AV_INPUT_BUFFER_PADDING_SIZE);
94            if (!filtered_buf) {
95                return AVERROR(ENOMEM);
96            }
97            memset(filtered_buf->data + filtered_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
98        }
99
100        extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
101        if (!extradata) {
102            av_buffer_unref(&filtered_buf);
103            return AVERROR(ENOMEM);
104        }
105
106        *data = extradata;
107        *size = extradata_size;
108
109        bytestream2_init_writer(&pb_extradata, extradata, extradata_size);
110        if (s->remove)
111            bytestream2_init_writer(&pb_filtered_data, filtered_buf->data, filtered_size);
112
113        for (i = 0; i < s->av1_pkt.nb_obus; i++) {
114            AV1OBU *obu = &s->av1_pkt.obus[i];
115            if (val_in_array(extradata_obu_types, nb_extradata_obu_types,
116                             obu->type)) {
117                bytestream2_put_bufferu(&pb_extradata, obu->raw_data, obu->raw_size);
118            } else if (s->remove) {
119                bytestream2_put_bufferu(&pb_filtered_data, obu->raw_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
134static 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, 0);
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 = NULL;
183        PutByteContext pb_filtered_data, pb_extradata;
184        uint8_t *extradata;
185
186        if (s->remove) {
187            filtered_buf = av_buffer_alloc(filtered_size + AV_INPUT_BUFFER_PADDING_SIZE);
188            if (!filtered_buf) {
189                return AVERROR(ENOMEM);
190            }
191            memset(filtered_buf->data + filtered_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
192        }
193
194        extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
195        if (!extradata) {
196            av_buffer_unref(&filtered_buf);
197            return AVERROR(ENOMEM);
198        }
199
200        *data = extradata;
201        *size = extradata_size;
202
203        bytestream2_init_writer(&pb_extradata, extradata, extradata_size);
204        if (s->remove)
205            bytestream2_init_writer(&pb_filtered_data, filtered_buf->data, filtered_size);
206
207        for (i = 0; i < s->h2645_pkt.nb_nals; i++) {
208            H2645NAL *nal = &s->h2645_pkt.nals[i];
209            if (val_in_array(extradata_nal_types, nb_extradata_nal_types,
210                             nal->type)) {
211                bytestream2_put_be24u(&pb_extradata, 1); //startcode
212                bytestream2_put_bufferu(&pb_extradata, nal->raw_data, nal->raw_size);
213            } else if (s->remove) {
214                bytestream2_put_be24u(&pb_filtered_data, 1); // startcode
215                bytestream2_put_bufferu(&pb_filtered_data, nal->raw_data, 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
230static 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        *size = extradata_size;
255
256        if (s->remove) {
257            pkt->data += extradata_size;
258            pkt->size -= extradata_size;
259        }
260    }
261
262    return 0;
263}
264
265static int extract_extradata_mpeg12(AVBSFContext *ctx, AVPacket *pkt,
266                                     uint8_t **data, int *size)
267{
268    ExtractExtradataContext *s = ctx->priv_data;
269    uint32_t state = UINT32_MAX;
270    int i, found = 0;
271
272    for (i = 0; i < pkt->size; i++) {
273        state = (state << 8) | pkt->data[i];
274        if (state == 0x1B3)
275            found = 1;
276        else if (found && state != 0x1B5 && state < 0x200 && state >= 0x100) {
277            *size = i - 3;
278            *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
279            if (!*data)
280                return AVERROR(ENOMEM);
281
282            memcpy(*data, pkt->data, *size);
283
284            if (s->remove) {
285                pkt->data += *size;
286                pkt->size -= *size;
287            }
288            break;
289        }
290    }
291    return 0;
292}
293
294static int extract_extradata_mpeg4(AVBSFContext *ctx, AVPacket *pkt,
295                                   uint8_t **data, int *size)
296{
297    ExtractExtradataContext *s = ctx->priv_data;
298    const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size;
299    uint32_t state = UINT32_MAX;
300
301    while (ptr < end) {
302        ptr = avpriv_find_start_code(ptr, end, &state);
303        if (state == 0x1B3 || state == 0x1B6) {
304            if (ptr - pkt->data > 4) {
305                *size = ptr - 4 - pkt->data;
306                *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
307                if (!*data)
308                    return AVERROR(ENOMEM);
309
310                memcpy(*data, pkt->data, *size);
311
312                if (s->remove) {
313                    pkt->data += *size;
314                    pkt->size -= *size;
315                }
316            }
317            break;
318        }
319    }
320    return 0;
321}
322
323static const struct {
324    enum AVCodecID id;
325    int (*extract)(AVBSFContext *ctx, AVPacket *pkt,
326                   uint8_t **data, int *size);
327} extract_tab[] = {
328    { AV_CODEC_ID_AV1,        extract_extradata_av1     },
329    { AV_CODEC_ID_AVS2,       extract_extradata_mpeg4   },
330    { AV_CODEC_ID_AVS3,       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
340static 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
357static 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
384fail:
385    av_packet_unref(pkt);
386    return ret;
387}
388
389static 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
396static const enum AVCodecID codec_ids[] = {
397    AV_CODEC_ID_AV1,
398    AV_CODEC_ID_AVS2,
399    AV_CODEC_ID_AVS3,
400    AV_CODEC_ID_CAVS,
401    AV_CODEC_ID_H264,
402    AV_CODEC_ID_HEVC,
403    AV_CODEC_ID_MPEG1VIDEO,
404    AV_CODEC_ID_MPEG2VIDEO,
405    AV_CODEC_ID_MPEG4,
406    AV_CODEC_ID_VC1,
407    AV_CODEC_ID_NONE,
408};
409
410#define OFFSET(x) offsetof(ExtractExtradataContext, x)
411#define FLAGS (AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_BSF_PARAM)
412static const AVOption options[] = {
413    { "remove", "remove the extradata from the bitstream", OFFSET(remove), AV_OPT_TYPE_INT,
414        { .i64 = 0 }, 0, 1, FLAGS },
415    { NULL },
416};
417
418static const AVClass extract_extradata_class = {
419    .class_name = "extract_extradata",
420    .item_name  = av_default_item_name,
421    .option     = options,
422    .version    = LIBAVUTIL_VERSION_INT,
423};
424
425const FFBitStreamFilter ff_extract_extradata_bsf = {
426    .p.name         = "extract_extradata",
427    .p.codec_ids    = codec_ids,
428    .p.priv_class   = &extract_extradata_class,
429    .priv_data_size = sizeof(ExtractExtradataContext),
430    .init           = extract_extradata_init,
431    .filter         = extract_extradata_filter,
432    .close          = extract_extradata_close,
433};
434