[FFmpeg-devel] [PATCH v4] avcodec/libx264: fix extradata when config annexb=0

[Andreas Rheinhardt]

Zhao Zhili:
> From: Zhao Zhili <zhilizhao at tencent.com>
> 
> ---
> v4: Fix missing SEI in set_avcc_extradata
> v3: Remove unnecessary inclusion
> 
>  configure            |   2 +-
>  libavcodec/libx264.c | 153 ++++++++++++++++++++++++++++++++++++-------
>  2 files changed, 130 insertions(+), 25 deletions(-)
> 
> diff --git a/configure b/configure
> index db7dc89755..24cb897d28 100755
> --- a/configure
> +++ b/configure
> @@ -3491,7 +3491,7 @@ libwebp_encoder_deps="libwebp"
>  libwebp_anim_encoder_deps="libwebp"
>  libx262_encoder_deps="libx262"
>  libx264_encoder_deps="libx264"
> -libx264_encoder_select="atsc_a53"
> +libx264_encoder_select="atsc_a53 h264parse"
>  libx264rgb_encoder_deps="libx264"
>  libx264rgb_encoder_select="libx264_encoder"
>  libx265_encoder_deps="libx265"
> diff --git a/libavcodec/libx264.c b/libavcodec/libx264.c
> index 10d646bd76..e7d16997d2 100644
> --- a/libavcodec/libx264.c
> +++ b/libavcodec/libx264.c
> @@ -34,6 +34,7 @@
>  #include "avcodec.h"
>  #include "codec_internal.h"
>  #include "encode.h"
> +#include "h264_ps.h"
>  #include "internal.h"
>  #include "packet_internal.h"
>  #include "atsc_a53.h"
> @@ -865,6 +866,131 @@ static int convert_pix_fmt(enum AVPixelFormat pix_fmt)
>      return 0;
>  }
>  
> +static int save_sei(AVCodecContext *avctx, x264_nal_t *nal)
> +{
> +    X264Context *x4 = avctx->priv_data;
> +
> +    av_log(avctx, AV_LOG_INFO, "%s\n", nal->p_payload + 25);
> +    x4->sei_size = nal->i_payload;
> +    x4->sei = av_malloc(x4->sei_size);
> +    if (!x4->sei)
> +        return AVERROR(ENOMEM);
> +
> +    memcpy(x4->sei, nal->p_payload, nal->i_payload);
> +
> +    return 0;
> +}
> +
> +static int set_avcc_extradata(AVCodecContext *avctx, x264_nal_t *nal, int nnal)
> +{
> +    X264Context *x4 = avctx->priv_data;
> +    x264_nal_t *sps_nal = NULL;
> +    x264_nal_t *pps_nal = NULL;
> +    uint8_t *p, *sps;
> +    int ret;
> +
> +    /* We know it's in the order of SPS/PPS/SEI, but it's not documented in x264 API.
> +     * The x264 param i_sps_id implies there is a single pair of SPS/PPS.
> +     */
> +    for (int i = 0; i < nnal; i++) {
> +        switch (nal[i].i_type) {
> +        case NAL_SPS:
> +            sps_nal = &nal[i];
> +            break;
> +        case NAL_PPS:
> +            pps_nal = &nal[i];
> +            break;
> +        case NAL_SEI:
> +            ret = save_sei(avctx, &nal[i]);
> +            if (ret < 0)
> +                return ret;
> +            break;
> +        }
> +    }
> +    if (!sps_nal || !pps_nal)
> +        return AVERROR_EXTERNAL;
> +
> +    avctx->extradata_size = sps_nal->i_payload + pps_nal->i_payload + 7;
> +    avctx->extradata = av_mallocz(avctx->extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
> +    if (!avctx->extradata)
> +        return AVERROR(ENOMEM);
> +
> +    // Now create AVCDecoderConfigurationRecord
> +    p = avctx->extradata;
> +    // Skip size part
> +    sps = sps_nal->p_payload + 4;
> +    *p++ = 1; // version
> +    *p++ = sps[1]; // AVCProfileIndication
> +    *p++ = sps[2]; // profile_compatibility
> +    *p++ = sps[3]; // AVCLevelIndication
> +    *p++ = 0xFF;
> +    *p++ = 0xE0 | 0x01; // 3 bits reserved (111) + 5 bits number of sps
> +    memcpy(p, sps_nal->p_payload + 2, sps_nal->i_payload - 2);
> +    // Make sps has AV_INPUT_BUFFER_PADDING_SIZE padding, so it can be used
> +    // with GetBitContext
> +    sps = p + 2;
> +    p += sps_nal->i_payload - 2;
> +    *p++ = 1;
> +    memcpy(p, pps_nal->p_payload + 2, pps_nal->i_payload - 2);
> +    p += pps_nal->i_payload - 2;
> +
> +    if (sps[3] != 66 && sps[3] != 77 && sps[3] != 88) {
> +        GetBitContext gbc;
> +        H264ParamSets ps = { 0 };
> +
> +        init_get_bits8(&gbc, sps, sps_nal->i_payload - 4);
> +        skip_bits(&gbc, 8);
> +        ret = ff_h264_decode_seq_parameter_set(&gbc, avctx, &ps, 1);

ff_h264_decode_seq_parameter_set() expects to read from a GetBitContext
whose buffer has already been stripped of 0x03 escape bytes. Your buffer
hasn't and therefore it is possible for this function to return an error
even when the input is fine.

> +        if (ret < 0)
> +            return ret;
> +
> +        ps.sps = ps.sps_list[x4->params.i_sps_id];
> +        *p++ = 0xFC | ps.sps->chroma_format_idc;
> +        *p++ = 0xF8 | (ps.sps->bit_depth_luma - 8);
> +        *p++ = 0xF8 | (ps.sps->bit_depth_chroma - 8);
> +        *p++ = 0;
> +        ff_h264_ps_uninit(&ps);
> +    }
> +    av_assert0(avctx->extradata + avctx->extradata_size >= p);
> +    avctx->extradata_size = p - avctx->extradata;
> +
> +    return 0;
> +}
> +
> +static int set_extradata(AVCodecContext *avctx)
> +{
> +    X264Context *x4 = avctx->priv_data;
> +    x264_nal_t *nal;
> +    uint8_t *p;
> +    int nnal, s;
> +
> +    s = x264_encoder_headers(x4->enc, &nal, &nnal);
> +    if (s < 0)
> +        return AVERROR_EXTERNAL;
> +
> +    if (!x4->params.b_annexb)
> +        return set_avcc_extradata(avctx, nal, nnal);
> +
> +    avctx->extradata = p = av_mallocz(s + AV_INPUT_BUFFER_PADDING_SIZE);
> +    if (!p)
> +        return AVERROR(ENOMEM);
> +
> +    for (int i = 0; i < nnal; i++) {
> +        /* Don't put the SEI in extradata. */
> +        if (nal[i].i_type == NAL_SEI) {
> +            s = save_sei(avctx, &nal[i]);
> +            if (s < 0)
> +                return s;
> +            continue;
> +        }
> +        memcpy(p, nal[i].p_payload, nal[i].i_payload);
> +        p += nal[i].i_payload;
> +    }
> +    avctx->extradata_size = p - avctx->extradata;
> +
> +    return 0;
> +}
> +
>  #define PARSE_X264_OPT(name, var)\
>      if (x4->var && x264_param_parse(&x4->params, name, x4->var) < 0) {\
>          av_log(avctx, AV_LOG_ERROR, "Error parsing option '%s' with value '%s'.\n", name, x4->var);\
> @@ -1233,30 +1359,9 @@ FF_ENABLE_DEPRECATION_WARNINGS
>          return AVERROR_EXTERNAL;
>  
>      if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
> -        x264_nal_t *nal;
> -        uint8_t *p;
> -        int nnal, s, i;
> -
> -        s = x264_encoder_headers(x4->enc, &nal, &nnal);
> -        avctx->extradata = p = av_mallocz(s + AV_INPUT_BUFFER_PADDING_SIZE);
> -        if (!p)
> -            return AVERROR(ENOMEM);
> -
> -        for (i = 0; i < nnal; i++) {
> -            /* Don't put the SEI in extradata. */
> -            if (nal[i].i_type == NAL_SEI) {
> -                av_log(avctx, AV_LOG_INFO, "%s\n", nal[i].p_payload+25);
> -                x4->sei_size = nal[i].i_payload;
> -                x4->sei      = av_malloc(x4->sei_size);
> -                if (!x4->sei)
> -                    return AVERROR(ENOMEM);
> -                memcpy(x4->sei, nal[i].p_payload, nal[i].i_payload);
> -                continue;
> -            }
> -            memcpy(p, nal[i].p_payload, nal[i].i_payload);
> -            p += nal[i].i_payload;
> -        }
> -        avctx->extradata_size = p - avctx->extradata;
> +        ret = set_extradata(avctx);
> +        if (ret < 0)
> +            return ret;
>      }
>  
>      cpb_props = ff_encode_add_cpb_side_data(avctx);