[FFmpeg-devel] [PATCH 3/6] avcodec/extract_extradata_bsf: use the parsing code from h264_split()

[James Almer]

The current code discards every non SPS/PPS/VPS NAL from extradata.
No samples where such filtering breaks decoding has been found, but it's
nonetheless safer to keep the current behaviour from AVCodecParser for the
time being, or until the failures with h264 are fixed.

Signed-off-by: James Almer <jamrial at gmail.com>
---
 libavcodec/extract_extradata_bsf.c | 107 ++++++++++++-------------------------
 1 file changed, 35 insertions(+), 72 deletions(-)

diff --git a/libavcodec/extract_extradata_bsf.c b/libavcodec/extract_extradata_bsf.c
index bb2e9bf68c..89981996f2 100644
--- a/libavcodec/extract_extradata_bsf.c
+++ b/libavcodec/extract_extradata_bsf.c
@@ -40,15 +40,6 @@ typedef struct ExtractExtradataContext {
     int remove;
 } ExtractExtradataContext;
 
-static int val_in_array(const int *arr, int len, int val)
-{
-    int i;
-    for (i = 0; i < len; i++)
-        if (arr[i] == val)
-            return 1;
-    return 0;
-}
-
 static int extract_extradata_h264(AVBSFContext *ctx, AVPacket *pkt,
                                   uint8_t **data, int *size)
 {
@@ -98,77 +89,49 @@ static int extract_extradata_h264(AVBSFContext *ctx, AVPacket *pkt,
 static int extract_extradata_h265(AVBSFContext *ctx, AVPacket *pkt,
                                   uint8_t **data, int *size)
 {
-    static const int extradata_nal_types_hevc[] = {
-        HEVC_NAL_VPS, HEVC_NAL_SPS, HEVC_NAL_PPS,
-    };
-
     ExtractExtradataContext *s = ctx->priv_data;
+    uint32_t state = UINT32_MAX;
+    const uint8_t *ptr = pkt->data, *end = pkt->data + pkt->size;
+    int has_vps = 0;
+    int has_sps = 0;
+    int has_pps = 0;
+    int nalu_type;
 
-    H2645Packet h2645_pkt = { 0 };
-    int extradata_size = 0;
-    const int *extradata_nal_types;
-    int nb_extradata_nal_types;
-    int i, ret = 0;
-
-    extradata_nal_types    = extradata_nal_types_hevc;
-    nb_extradata_nal_types = FF_ARRAY_ELEMS(extradata_nal_types_hevc);
-
-    ret = ff_h2645_packet_split(&h2645_pkt, pkt->data, pkt->size,
-                                ctx, 0, 0, ctx->par_in->codec_id, 1);
-    if (ret < 0)
-        return ret;
-
-    for (i = 0; i < h2645_pkt.nb_nals; i++) {
-        H2645NAL *nal = &h2645_pkt.nals[i];
-        if (val_in_array(extradata_nal_types, nb_extradata_nal_types, nal->type))
-            extradata_size += nal->raw_size + 3;
-    }
-
-    if (extradata_size) {
-        AVBufferRef *filtered_buf;
-        uint8_t *extradata, *filtered_data;
-
-        if (s->remove) {
-            filtered_buf = av_buffer_alloc(pkt->size + AV_INPUT_BUFFER_PADDING_SIZE);
-            if (!filtered_buf)
-                goto fail;
-            filtered_data = filtered_buf->data;
-        }
-
-        extradata = av_malloc(extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
-        if (!extradata) {
-            av_buffer_unref(&filtered_buf);
-            goto fail;
-        }
+    while (ptr < end) {
+        ptr = avpriv_find_start_code(ptr, end, &state);
+        if ((state & 0xFFFFFF00) != 0x100)
+            break;
+        nalu_type = (state >> 1) & 0x3F;
+        if (nalu_type == HEVC_NAL_VPS)
+            has_vps = 1;
+        else if (nalu_type == HEVC_NAL_SPS)
+            has_sps = 1;
+        else if (nalu_type == HEVC_NAL_PPS)
+            has_pps = 1;
+        else if ((nalu_type != HEVC_NAL_SEI_PREFIX || has_pps) &&
+                  nalu_type != HEVC_NAL_AUD) {
+            if (has_vps && has_sps) {
+                while (ptr - 4 > pkt->data && ptr[-5] == 0)
+                    ptr--;
+                if (ptr - pkt->data > 4) {
+                    *size = ptr - 4 - pkt->data;
+                    *data = av_malloc(*size + AV_INPUT_BUFFER_PADDING_SIZE);
+                    if (!*data)
+                        return AVERROR(ENOMEM);
 
-        *data = extradata;
-        *size = extradata_size;
+                    memcpy(*data, pkt->data, *size);
 
-        for (i = 0; i < h2645_pkt.nb_nals; i++) {
-            H2645NAL *nal = &h2645_pkt.nals[i];
-            if (val_in_array(extradata_nal_types, nb_extradata_nal_types,
-                             nal->type)) {
-                AV_WB24(extradata, 1); // startcode
-                memcpy(extradata + 3, nal->raw_data, nal->raw_size);
-                extradata += 3 + nal->raw_size;
-            } else if (s->remove) {
-                AV_WB24(filtered_data, 1); // startcode
-                memcpy(filtered_data + 3, nal->raw_data, nal->raw_size);
-                filtered_data += 3 + nal->raw_size;
+                    if (s->remove) {
+                        pkt->data += *size;
+                        pkt->size -= *size;
+                    }
+                }
+                return 0;
             }
         }
-
-        if (s->remove) {
-            av_buffer_unref(&pkt->buf);
-            pkt->buf  = filtered_buf;
-            pkt->data = filtered_buf->data;
-            pkt->size = filtered_data - filtered_buf->data;
-        }
     }
 
-fail:
-    ff_h2645_packet_uninit(&h2645_pkt);
-    return ret;
+    return 0;
 }
 
 static int extract_extradata_vc1(AVBSFContext *ctx, AVPacket *pkt,
-- 
2.12.0