diff mbox series

[FFmpeg-devel,16/17] avformat/avc: Improve writing AVCDecoderConfigurationRecord

Message ID 20200709192022.9412-8-andreas.rheinhardt@gmail.com
State New
Headers show
Series [FFmpeg-devel,1/7] avformat/avc: Fix undefined shift and assert when reading exp-golomb num
Related show

Checks

Context Check Description
andriy/default pending
andriy/make success Make finished
andriy/make_fate success Make fate finished

Commit Message

Andreas Rheinhardt July 9, 2020, 7:20 p.m. UTC
According to the specifications, both the SPS as well as the PPS in
an AVCDecoderConfigurationRecord (the mp4-style extradata for H.264) are
supposed to be ordered according to increasing SPS/PPS id.
ff_isom_write_avcc did not ensure this; in fact, it did not even ensure
that there are no two parameter sets of the same type with the same id.

This commit changes this. Only the last occurence of any SPS/PPS id is
kept, i.e. just like with in-band parameter sets, new parameter sets
overwrite old parameter sets with the same id, and those parameter sets
that are kept are written in the order of their id.

This also led to the removal of one round of copying implicit in the
usage of dynamic buffers.

(mkvextract simply converts the mp4-style extradata contained in the
CodecPrivate of a Matroska file and prepends it to the frame data (which
is also converted to annex B); if the file had in-band extradata, the
created file would have the extradata twice at the beginning and
extract_extradata would include all of these parameter sets.)

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt@gmail.com>
---
 libavformat/avc.c | 131 +++++++++++++++++++++++-----------------------
 1 file changed, 66 insertions(+), 65 deletions(-)
diff mbox series

Patch

diff --git a/libavformat/avc.c b/libavformat/avc.c
index c106fe5792..0f43b295a4 100644
--- a/libavformat/avc.c
+++ b/libavformat/avc.c
@@ -167,13 +167,32 @@  int ff_avc_parse_nal_units_buf(const uint8_t *buf_in, uint8_t **buf, int *size)
     return 0;
 }
 
+static void write_array(AVIOContext *pb, const uint8_t *const *ps,
+                        const uint16_t *const ps_sizes,
+                        int nb_ps, int array_size, uint8_t reserved)
+{
+    avio_w8(pb, reserved | nb_ps);
+    for (int i = 0; i < array_size; i++) {
+        if (ps_sizes[i]) {
+            avio_wb16(pb, ps_sizes[i]);
+            avio_write(pb, ps[i], ps_sizes[i]);
+            nb_ps--;
+        }
+    }
+    av_assert1(nb_ps == 0);
+}
+
 int ff_isom_write_avcc(AVIOContext *pb, const uint8_t *data, int len)
 {
-    AVIOContext *sps_pb = NULL, *pps_pb = NULL, *sps_ext_pb = NULL;
     const uint8_t *nal, *nal_end, *end;
-    uint8_t *sps, *pps, *sps_ext;
-    uint32_t sps_size = 0, pps_size = 0, sps_ext_size = 0;
     int ret, nb_sps = 0, nb_pps = 0, nb_sps_ext = 0;
+    const uint8_t *sps[H264_MAX_SPS_COUNT];
+    const uint8_t *pps[H264_MAX_PPS_COUNT];
+    const uint8_t *sps_ext[H264_MAX_SPS_COUNT];
+    uint16_t sps_sizes[H264_MAX_SPS_COUNT]     = { 0 };
+    uint16_t pps_sizes[H264_MAX_PPS_COUNT]     = { 0 };
+    uint16_t sps_ext_sizes[H264_MAX_SPS_COUNT] = { 0 };
+    H264SPS seq;
 
     if (len <= 6)
         return AVERROR_INVALIDDATA;
@@ -185,16 +204,6 @@  int ff_isom_write_avcc(AVIOContext *pb, const uint8_t *data, int len)
         return 0;
     }
 
-    ret = avio_open_dyn_buf(&sps_pb);
-    if (ret < 0)
-        goto fail;
-    ret = avio_open_dyn_buf(&pps_pb);
-    if (ret < 0)
-        goto fail;
-    ret = avio_open_dyn_buf(&sps_ext_pb);
-    if (ret < 0)
-        goto fail;
-
     /* look for sps and pps */
     nal_end = NULL;
     end     = data + len;
@@ -203,71 +212,63 @@  int ff_isom_write_avcc(AVIOContext *pb, const uint8_t *data, int len)
         uint8_t nal_type = nal[0] & 0x1f;
 
         if (nal_type == 7) { /* SPS */
-            nb_sps++;
-            if (size > UINT16_MAX || nb_sps >= H264_MAX_SPS_COUNT) {
-                ret = AVERROR_INVALIDDATA;
-                goto fail;
-            }
-            avio_wb16(sps_pb, size);
-            avio_write(sps_pb, nal, size);
-        } else if (nal_type == 8) { /* PPS */
-            nb_pps++;
-            if (size > UINT16_MAX || nb_pps >= H264_MAX_PPS_COUNT) {
-                ret = AVERROR_INVALIDDATA;
-                goto fail;
-            }
-            avio_wb16(pps_pb, size);
-            avio_write(pps_pb, nal, size);
-        } else if (nal_type == 13) { /* SPS_EXT */
-            nb_sps_ext++;
-            if (size > UINT16_MAX || nb_sps_ext >= 256) {
-                ret = AVERROR_INVALIDDATA;
-                goto fail;
+            if (size > UINT16_MAX)
+                return AVERROR_INVALIDDATA;
+            ret = ff_avc_decode_sps(&seq, nal + 1, size - 1);
+            if (ret < 0)
+                return ret;
+            nb_sps           += 1 - !!sps_sizes[seq.id];
+            sps[seq.id]       = nal;
+            sps_sizes[seq.id] = size;
+        } else if (nal_type == H264_NAL_PPS || nal_type == H264_NAL_SPS_EXT) {
+            GetBitContext gb;
+            unsigned id;
+            if (size > UINT16_MAX)
+                return AVERROR_INVALIDDATA;
+            /* No need to unescape here */
+            ret = init_get_bits8(&gb, nal + 1, size - 1);
+            av_assert1(ret >= 0);
+            ret = get_ue_golomb(&gb, &id);
+            if (ret < 0)
+                return ret;
+            if (nal_type == H264_NAL_PPS) {
+                if (id >= H264_MAX_PPS_COUNT)
+                    return AVERROR_INVALIDDATA;
+                nb_pps       += 1 - !!pps_sizes[id];
+                pps[id]       = nal;
+                pps_sizes[id] = size;
+            } else {
+                if (id >= H264_MAX_SPS_COUNT)
+                    return AVERROR_INVALIDDATA;
+                if (!sps_sizes[id])
+                    continue;
+                nb_sps_ext       += 1 - !!sps_ext_sizes[id];
+                sps_ext[id]       = nal;
+                sps_ext_sizes[id] = size;
             }
-            avio_wb16(sps_ext_pb, size);
-            avio_write(sps_ext_pb, nal, size);
         }
     }
-    sps_size = avio_get_dyn_buf(sps_pb, &sps);
-    pps_size = avio_get_dyn_buf(pps_pb, &pps);
-    sps_ext_size = avio_get_dyn_buf(sps_ext_pb, &sps_ext);
 
-    if (sps_size < 6 || !pps_size) {
-        ret = AVERROR_INVALIDDATA;
-        goto fail;
-    }
+    if (!nb_sps || nb_sps == H264_MAX_SPS_COUNT ||
+        !nb_pps || nb_pps == H264_MAX_PPS_COUNT)
+        return AVERROR_INVALIDDATA;
 
     avio_w8(pb, 1); /* version */
-    avio_w8(pb, sps[3]); /* profile */
-    avio_w8(pb, sps[4]); /* profile compat */
-    avio_w8(pb, sps[5]); /* level */
+    avio_w8(pb, seq.profile_idc);
+    avio_w8(pb, seq.constraint_set_flags); /* profile compat */
+    avio_w8(pb, seq.level_idc);
     avio_w8(pb, 0xff); /* 6 bits reserved (111111) + 2 bits nal size length - 1 (11) */
-    avio_w8(pb, 0xe0 | nb_sps); /* 3 bits reserved (111) + 5 bits number of sps */
-
-    avio_write(pb, sps, sps_size);
-    avio_w8(pb, nb_pps); /* number of pps */
-    avio_write(pb, pps, pps_size);
-
-    if (sps[3] != 66 && sps[3] != 77 && sps[3] != 88) {
-        H264SPS seq;
-        ret = ff_avc_decode_sps(&seq, sps + 3, sps_size - 3);
-        if (ret < 0)
-            goto fail;
+    write_array(pb, sps, sps_sizes, nb_sps, H264_MAX_SPS_COUNT, 0xe0);
+    write_array(pb, pps, pps_sizes, nb_pps, H264_MAX_PPS_COUNT, 0x00);
 
+    if (seq.profile_idc != 66 && seq.profile_idc != 77 && seq.profile_idc != 88) {
         avio_w8(pb, 0xfc |  seq.chroma_format_idc); /* 6 bits reserved (111111) + chroma_format_idc */
         avio_w8(pb, 0xf8 | seq.bit_depth_luma_minus8); /* 5 bits reserved (11111) + bit_depth_luma_minus8 */
         avio_w8(pb, 0xf8 | seq.bit_depth_chroma_minus8); /* 5 bits reserved (11111) + bit_depth_chroma_minus8 */
-        avio_w8(pb, nb_sps_ext); /* number of sps ext */
-        if (nb_sps_ext)
-            avio_write(pb, sps_ext, sps_ext_size);
+        write_array(pb, sps_ext, sps_ext_sizes, nb_sps_ext, H264_MAX_SPS_COUNT, 0);
     }
 
-fail:
-    ffio_free_dyn_buf(&sps_pb);
-    ffio_free_dyn_buf(&pps_pb);
-    ffio_free_dyn_buf(&sps_ext_pb);
-
-    return ret;
+    return 0;
 }
 
 int ff_avc_write_annexb_extradata(const uint8_t *in, uint8_t **buf, int *size)