[FFmpeg-devel] [PATCH v2 2/9] cbs_vp9: Implement parser entrypoint

[Mark Thompson]

---
 libavcodec/cbs_vp9.c | 90 +++++++++++++++++++++++++++++++++++---------
 1 file changed, 73 insertions(+), 17 deletions(-)

diff --git a/libavcodec/cbs_vp9.c b/libavcodec/cbs_vp9.c
index 0b5f137ed8..6ea4681d68 100644
--- a/libavcodec/cbs_vp9.c
+++ b/libavcodec/cbs_vp9.c
@@ -409,15 +409,23 @@ static int cbs_vp9_write_le(CodedBitstreamContext *ctx, PutBitContext *pbc,
 #undef byte_alignment
 
 
-static int cbs_vp9_split_fragment(CodedBitstreamContext *ctx,
-                                  CodedBitstreamFragment *frag,
-                                  int header)
+typedef int (*cbs_vp9_split_frame_callback)(CodedBitstreamContext *ctx,
+                                            void *priv,
+                                            const uint8_t *data,
+                                            size_t data_size);
+
+static int cbs_vp9_split_frames(CodedBitstreamContext *ctx,
+                                void *priv, cbs_vp9_split_frame_callback cb,
+                                const uint8_t *data, size_t data_size)
 {
     uint8_t superframe_header;
     int err;
 
+    if (data_size == 0)
+        return 0;
+
     // Last byte in the packet.
-    superframe_header = frag->data[frag->data_size - 1];
+    superframe_header = data[data_size - 1];
 
     if ((superframe_header & 0xe0) == 0xc0) {
         VP9RawSuperframeIndex sfi;
@@ -427,8 +435,14 @@ static int cbs_vp9_split_fragment(CodedBitstreamContext *ctx,
 
         index_size = 2 + (((superframe_header & 0x18) >> 3) + 1) *
                           ((superframe_header & 0x07) + 1);
+        if (index_size > data_size) {
+            av_log(ctx->log_ctx, AV_LOG_ERROR, "Superframe index (%"
+                   SIZE_SPECIFIER" bytes) is larger than whole frame (%"
+                   SIZE_SPECIFIER" bytes).\n", index_size, data_size);
+            return AVERROR_INVALIDDATA;
+        }
 
-        err = init_get_bits(&gbc, frag->data + frag->data_size - index_size,
+        err = init_get_bits(&gbc, data + data_size - index_size,
                             8 * index_size);
         if (err < 0)
             return err;
@@ -439,34 +453,27 @@ static int cbs_vp9_split_fragment(CodedBitstreamContext *ctx,
 
         pos = 0;
         for (i = 0; i <= sfi.frames_in_superframe_minus_1; i++) {
-            if (pos + sfi.frame_sizes[i] + index_size > frag->data_size) {
+            if (pos + sfi.frame_sizes[i] + index_size > data_size) {
                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Frame %d too large "
                        "in superframe: %"PRIu32" bytes.\n",
                        i, sfi.frame_sizes[i]);
                 return AVERROR_INVALIDDATA;
             }
 
-            err = ff_cbs_insert_unit_data(ctx, frag, -1, 0,
-                                          frag->data + pos,
-                                          sfi.frame_sizes[i],
-                                          frag->data_ref);
+            err = cb(ctx, priv, data + pos, sfi.frame_sizes[i]);
             if (err < 0)
                 return err;
 
             pos += sfi.frame_sizes[i];
         }
-        if (pos + index_size != frag->data_size) {
+        if (pos + index_size != data_size) {
             av_log(ctx->log_ctx, AV_LOG_WARNING, "Extra padding at "
                    "end of superframe: %"SIZE_SPECIFIER" bytes.\n",
-                   frag->data_size - (pos + index_size));
+                   data_size - (pos + index_size));
         }
 
-        return 0;
-
     } else {
-        err = ff_cbs_insert_unit_data(ctx, frag, -1, 0,
-                                      frag->data, frag->data_size,
-                                      frag->data_ref);
+        err = cb(ctx, priv, data, data_size);
         if (err < 0)
             return err;
     }
@@ -474,6 +481,23 @@ static int cbs_vp9_split_fragment(CodedBitstreamContext *ctx,
     return 0;
 }
 
+static int cbs_vp9_insert_unit(CodedBitstreamContext *ctx, void *priv,
+                               const uint8_t *data, size_t data_size)
+{
+    CodedBitstreamFragment *frag = priv;
+    return ff_cbs_insert_unit_data(ctx, frag, -1, 0,
+                                   (uint8_t*)data, data_size,
+                                   frag->data_ref);
+}
+
+static int cbs_vp9_split_fragment(CodedBitstreamContext *ctx,
+                                  CodedBitstreamFragment *frag,
+                                  int header)
+{
+    return cbs_vp9_split_frames(ctx, frag, &cbs_vp9_insert_unit,
+                                frag->data, frag->data_size);
+}
+
 static void cbs_vp9_free_frame(void *unit, uint8_t *content)
 {
     VP9RawFrame *frame = (VP9RawFrame*)content;
@@ -678,6 +702,36 @@ static void cbs_vp9_close(CodedBitstreamContext *ctx)
     av_freep(&priv->write_buffer);
 }
 
+static int cbs_vp9_parse_frame(CodedBitstreamContext *ctx, void *priv,
+                               const uint8_t *data, size_t data_size)
+{
+    VP9RawFrameHeader *frame = priv;
+    GetBitContext gbc;
+    int err;
+
+    err = init_get_bits(&gbc, data, 8 * data_size);
+    if (err < 0)
+        return err;
+
+    memset(frame, 0, sizeof(*frame));
+    return cbs_vp9_read_uncompressed_header(ctx, &gbc, frame);
+}
+
+static int cbs_vp9_parse_headers(CodedBitstreamContext *ctx, void *header,
+                                 const uint8_t *data, size_t size)
+{
+    VP9RawFrameHeader frame;
+    int err;
+
+    err = cbs_vp9_split_frames(ctx, &frame, &cbs_vp9_parse_frame, data, size);
+    if (err < 0)
+        return err;
+
+    if (header)
+        memcpy(header, &frame, sizeof(frame));
+    return 0;
+}
+
 const CodedBitstreamType ff_cbs_type_vp9 = {
     .codec_id          = AV_CODEC_ID_VP9,
 
@@ -689,4 +743,6 @@ const CodedBitstreamType ff_cbs_type_vp9 = {
     .assemble_fragment = &cbs_vp9_assemble_fragment,
 
     .close             = &cbs_vp9_close,
+
+    .parse_headers     = &cbs_vp9_parse_headers,
 };
-- 
2.20.1