From 7bacf60ae5df75954a538563d19f6001aa598b3f Mon Sep 17 00:00:00 2001
From: Mark Thompson <sw@jkqxz.net>
Date: Mon, 23 Nov 2020 17:53:57 +0000
Subject: [PATCH] cbs: Add function to read extradata from an AVCodecContext

This is useful in decoders and parsers, matching the way that bitstream
filters read extradata from AVCodecParameters.
---
 libavcodec/cbs.c | 81 ++++++++++++++++++++++++------------------------
 libavcodec/cbs.h | 11 +++++++
 2 files changed, 52 insertions(+), 40 deletions(-)

diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
index c8c526ab12..f98531e131 100644
--- a/libavcodec/cbs.c
+++ b/libavcodec/cbs.c
@@ -223,66 +223,67 @@ static int cbs_fill_fragment_data(CodedBitstreamFragment *frag,
     return 0;
 }
 
-int ff_cbs_read_extradata(CodedBitstreamContext *ctx,
-                          CodedBitstreamFragment *frag,
-                          const AVCodecParameters *par)
+static int cbs_read_data(CodedBitstreamContext *ctx,
+                         CodedBitstreamFragment *frag,
+                         AVBufferRef *buf,
+                         const uint8_t *data, size_t size,
+                         int header)
 {
     int err;
 
-    err = cbs_fill_fragment_data(frag, par->extradata,
-                                 par->extradata_size);
-    if (err < 0)
-        return err;
+    if (buf) {
+        frag->data_ref = av_buffer_ref(buf);
+        if (!frag->data_ref)
+            return AVERROR(ENOMEM);
 
-    err = ctx->codec->split_fragment(ctx, frag, 1);
+        frag->data      = (uint8_t *)data;
+        frag->data_size = size;
+
+    } else {
+        err = cbs_fill_fragment_data(frag, data, size);
+        if (err < 0)
+            return err;
+    }
+
+    err = ctx->codec->split_fragment(ctx, frag, header);
     if (err < 0)
         return err;
 
     return cbs_read_fragment_content(ctx, frag);
 }
 
+int ff_cbs_read_extradata(CodedBitstreamContext *ctx,
+                          CodedBitstreamFragment *frag,
+                          const AVCodecParameters *par)
+{
+    return cbs_read_data(ctx, frag, NULL,
+                         par->extradata,
+                         par->extradata_size, 1);
+}
+
+int ff_cbs_read_extradata_from_codec(CodedBitstreamContext *ctx,
+                                     CodedBitstreamFragment *frag,
+                                     const AVCodecContext *avctx)
+{
+    return cbs_read_data(ctx, frag, NULL,
+                         avctx->extradata,
+                         avctx->extradata_size, 1);
+}
+
 int ff_cbs_read_packet(CodedBitstreamContext *ctx,
                        CodedBitstreamFragment *frag,
                        const AVPacket *pkt)
 {
-    int err;
-
-    if (pkt->buf) {
-        frag->data_ref = av_buffer_ref(pkt->buf);
-        if (!frag->data_ref)
-            return AVERROR(ENOMEM);
-
-        frag->data      = pkt->data;
-        frag->data_size = pkt->size;
-
-    } else {
-        err = cbs_fill_fragment_data(frag, pkt->data, pkt->size);
-        if (err < 0)
-            return err;
-    }
-
-    err = ctx->codec->split_fragment(ctx, frag, 0);
-    if (err < 0)
-        return err;
-
-    return cbs_read_fragment_content(ctx, frag);
+    return cbs_read_data(ctx, frag, pkt->buf,
+                         pkt->data, pkt->size, 0);
 }
 
 int ff_cbs_read(CodedBitstreamContext *ctx,
                 CodedBitstreamFragment *frag,
                 const uint8_t *data, size_t size)
 {
-    int err;
-
-    err = cbs_fill_fragment_data(frag, data, size);
-    if (err < 0)
-        return err;
-
-    err = ctx->codec->split_fragment(ctx, frag, 0);
-    if (err < 0)
-        return err;
-
-    return cbs_read_fragment_content(ctx, frag);
+    return cbs_read_data(ctx, frag, NULL,
+                         data, size, 0);
 }
 
 static int cbs_write_unit_data(CodedBitstreamContext *ctx,
diff --git a/libavcodec/cbs.h b/libavcodec/cbs.h
index 635921b11e..3fd0a0ef33 100644
--- a/libavcodec/cbs.h
+++ b/libavcodec/cbs.h
@@ -262,6 +262,17 @@ int ff_cbs_read_extradata(CodedBitstreamContext *ctx,
                           CodedBitstreamFragment *frag,
                           const AVCodecParameters *par);
 
+/**
+ * Read the extradata bitstream found in a codec context into a
+ * fragment, then split into units and decompose.
+ *
+ * This acts identical to ff_cbs_read_extradata() for the case where
+ * you already have a codec context.
+ */
+int ff_cbs_read_extradata_from_codec(CodedBitstreamContext *ctx,
+                                     CodedBitstreamFragment *frag,
+                                     const AVCodecContext *avctx);
+
 /**
  * Read the data bitstream from a packet into a fragment, then
  * split into units and decompose.