avformat/matroskaenc: Add API to write Masters with minimal length field

This muxer currently uses two ways to ensure that no bytes
are wasted by writing unnecessary long EBML length fields
for Master elements and the (Simple)Block element
(all the other elements are fine as one either already has
the right length or getting the actual length is easy
and necessary anyway):
Either use an upper bound that is good enough in case one
is available or write the data into a dynamic buffer first
to get the length; the former approach is impossible in
lots of cases, whereas the latter incurs allocations and
memcpying. It is therefore unfeasible to use the latter
for e.g. the attachments or the BlockGroups.
This patch adds a third alternative to complement the other two:
It consists of an EbmlWriter that one can add EBML elements to
that can be written later by calling ebml_writer_write();
the latter function first traverses the written elements recursively
and calculates the length of each element; then a second pass
is performed in which all the elements are written directly
(without any seeks).
This new API also performs checks for overlong elements;
this is in contrast to put_ebml_string() which simply performs
a size_t->int conversion even for strings originating from the user.

The new API is designed to have very low overhead: It uses
stack arrays and performs no allocations; this also comes
at a price: Right now, it can only be used in contexts in which
there is a compile-time upper bound for the number of elements.
It is also incompatible with storing the offset of an element
in order to update this field later. Furthermore, it puts
the onus of memory management (i.e. ensuring that pointers stay valid)
on the user.

These restrictions might be overcome in the future.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt@outlook.com>

1 files changed, 249 insertions, 0 deletions

diff --git a/libavformat/matroskaenc.c b/libavformat/matroskaenc.c
index 81194fd28d..086d403c85 100644
--- a/libavformat/matroskaenc.c
+++ b/libavformat/matroskaenc.c
@@ -61,6 +61,12 @@
  * Info, Tracks, Chapters, Attachments, Tags (potentially twice) and Cues */
 #define MAX_SEEKHEAD_ENTRIES 7
 
+/* Largest known-length EBML length */
+#define MAX_EBML_LENGTH ((1ULL << 56) - 2)
+/* The dynamic buffer API we rely upon has a limit of INT_MAX;
+ * and so has avio_write(). */
+#define MAX_SUPPORTED_EBML_LENGTH FFMIN(MAX_EBML_LENGTH, INT_MAX)
+
 #define MODE_MATROSKAv2 0x01
 #define MODE_WEBM       0x02
 
@@ -85,6 +91,48 @@ typedef struct ebml_stored_master {
     int64_t         pos;
 } ebml_stored_master;
 
+typedef enum EbmlType {
+    EBML_UINT,
+    EBML_SINT,
+    EBML_FLOAT,
+    EBML_UID,
+    EBML_STR,
+    EBML_UTF8 = EBML_STR,
+    EBML_BIN,
+    EBML_MASTER,
+} EbmlType;
+
+typedef struct EbmlMaster {
+    int nb_elements;               ///< -1 if not finished
+    int containing_master;         ///< -1 if no parent exists
+} EbmlMaster;
+
+typedef struct EbmlElement {
+    uint32_t id;
+    EbmlType type;
+    unsigned length_size;
+    uint64_t size;                 ///< excluding id and length field
+    union {
+        uint64_t uint;
+        int64_t  sint;
+        double   f;
+        const char    *str;
+        const uint8_t *bin;
+        EbmlMaster master;
+    } priv;
+} EbmlElement;
+
+typedef struct EbmlWriter {
+    unsigned nb_elements;
+    int      current_master_element;
+    EbmlElement *elements;
+} EbmlWriter;
+
+#define EBML_WRITER(max_nb_elems) \
+     EbmlElement elements[max_nb_elems]; \
+     EbmlWriter writer = (EbmlWriter){ .elements = elements, \
+                                       .current_master_element = -1 }
+
 typedef struct mkv_seekhead_entry {
     uint32_t        elementid;
     uint64_t        segmentpos;
@@ -362,6 +410,207 @@ static void end_ebml_master(AVIOContext *pb, ebml_master master)
     avio_seek(pb, pos, SEEK_SET);
 }
 
+static EbmlElement *ebml_writer_add(EbmlWriter *writer,
+                                    uint32_t id, EbmlType type)
+{
+    writer->elements[writer->nb_elements].id   = id;
+    writer->elements[writer->nb_elements].type = type;
+    return &writer->elements[writer->nb_elements++];
+}
+
+static void ebml_writer_open_master(EbmlWriter *writer, uint32_t id)
+{
+    EbmlElement *const elem = ebml_writer_add(writer, id, EBML_MASTER);
+    EbmlMaster *const master = &elem->priv.master;
+
+    master->containing_master = writer->current_master_element;
+    master->nb_elements = -1;
+
+    writer->current_master_element = writer->nb_elements - 1;
+}
+
+static void ebml_writer_add_string(EbmlWriter *writer, uint32_t id,
+                                   const char *str)
+{
+    EbmlElement *const elem = ebml_writer_add(writer, id, EBML_STR);
+
+    elem->priv.str = str;
+}
+
+static void ebml_writer_add_bin(EbmlWriter *writer, uint32_t id,
+                                const uint8_t *data, size_t size)
+{
+    EbmlElement *const elem = ebml_writer_add(writer, id, EBML_BIN);
+
+#if SIZE_MAX > UINT64_MAX
+    size = FFMIN(size, UINT64_MAX);
+#endif
+    elem->size = size;
+    elem->priv.bin = data;
+}
+
+static void ebml_writer_add_float(EbmlWriter *writer, uint32_t id,
+                                  double val)
+{
+    EbmlElement *const elem = ebml_writer_add(writer, id, EBML_FLOAT);
+
+    elem->priv.f = val;
+}
+
+static void ebml_writer_add_uid(EbmlWriter *writer, uint32_t id,
+                                uint64_t val)
+{
+    EbmlElement *const elem = ebml_writer_add(writer, id, EBML_UID);
+    elem->priv.uint = val;
+}
+
+static int ebml_writer_str_len(EbmlElement *elem)
+{
+    size_t len = strlen(elem->priv.str);
+#if SIZE_MAX > UINT64_MAX
+    len = FF_MIN(len, UINT64_MAX);
+#endif
+    elem->size = len;
+    return 0;
+}
+
+static av_const int uint_size(uint64_t val)
+{
+    int bytes = 0;
+    do {
+        bytes++;
+    } while (val >>= 8);
+    return bytes;
+}
+
+static int ebml_writer_uint_len(EbmlElement *elem)
+{
+    elem->size = uint_size(elem->priv.uint);
+    return 0;
+}
+
+static av_const int sint_size(int64_t val)
+{
+    uint64_t tmp = 2 * (uint64_t)(val < 0 ? val^-1 : val);
+    return uint_size(tmp);
+}
+
+static int ebml_writer_sint_len(EbmlElement *elem)
+{
+    elem->size = sint_size(elem->priv.sint);
+    return 0;
+}
+
+static int ebml_writer_elem_len(EbmlWriter *writer, EbmlElement *elem,
+                                int remaining_elems);
+
+static int ebml_writer_master_len(EbmlWriter *writer, EbmlElement *elem,
+                                  int remaining_elems)
+{
+    int nb_elems = elem->priv.master.nb_elements >= 0 ? elem->priv.master.nb_elements : remaining_elems - 1;
+    EbmlElement *const master = elem;
+    uint64_t total_size = 0;
+
+    master->priv.master.nb_elements = nb_elems;
+    for (; elem++, nb_elems > 0;) {
+        int ret = ebml_writer_elem_len(writer, elem, nb_elems);
+        if (ret < 0)
+            return ret;
+        av_assert2(ret < nb_elems);
+        /* No overflow is possible here, as both total_size and elem->size
+         * are bounded by MAX_SUPPORTED_EBML_LENGTH. */
+        total_size += ebml_id_size(elem->id) + elem->length_size + elem->size;
+        if (total_size > MAX_SUPPORTED_EBML_LENGTH)
+            return AVERROR(ERANGE);
+        nb_elems--;      /* consume elem */
+        elem += ret, nb_elems -= ret; /* and elem's children */
+    }
+    master->size = total_size;
+
+    return master->priv.master.nb_elements;
+}
+
+static int ebml_writer_elem_len(EbmlWriter *writer, EbmlElement *elem,
+                                int remaining_elems)
+{
+    int ret = 0;
+
+    switch (elem->type) {
+        case EBML_FLOAT:
+        case EBML_UID:
+            elem->size = 8;
+            break;
+        case EBML_STR:
+            ret = ebml_writer_str_len(elem);
+            break;
+        case EBML_UINT:
+            ret = ebml_writer_uint_len(elem);
+            break;
+        case EBML_SINT:
+            ret = ebml_writer_sint_len(elem);
+            break;
+        case EBML_MASTER:
+            ret = ebml_writer_master_len(writer, elem, remaining_elems);
+            break;
+    }
+    if (ret < 0)
+        return ret;
+    if (elem->size > MAX_SUPPORTED_EBML_LENGTH)
+        return AVERROR(ERANGE);
+    elem->length_size = ebml_length_size(elem->size);
+    return ret; /* number of elements consumed excluding elem itself */
+}
+
+static int ebml_writer_elem_write(const EbmlElement *elem, AVIOContext *pb)
+{
+    put_ebml_id(pb, elem->id);
+    put_ebml_num(pb, elem->size, elem->length_size);
+    switch (elem->type) {
+        case EBML_UID:
+        case EBML_FLOAT: {
+            uint64_t val = elem->type == EBML_UID ? elem->priv.uint
+                                                  : av_double2int(elem->priv.f);
+            avio_wb64(pb, val);
+            break;
+        }
+        case EBML_UINT:
+        case EBML_SINT: {
+            uint64_t val = elem->type == EBML_UINT ? elem->priv.uint
+                                                   : elem->priv.sint;
+            for (int i = elem->size; --i >= 0; )
+                avio_w8(pb, (uint8_t)(val >> i * 8));
+            break;
+        }
+        case EBML_STR:
+        case EBML_BIN: {
+            const uint8_t *data = elem->type == EBML_BIN ? elem->priv.bin
+                                         : (const uint8_t*)elem->priv.str;
+            avio_write(pb, data, elem->size);
+            break;
+        }
+        case EBML_MASTER: {
+            int nb_elems = elem->priv.master.nb_elements;
+
+            elem++;
+            for (int i = 0; i < nb_elems; i++)
+                i += ebml_writer_elem_write(elem + i, pb);
+
+            return nb_elems;
+        }
+    }
+    return 0;
+}
+
+static int ebml_writer_write(EbmlWriter *writer, AVIOContext *pb)
+{
+    int ret = ebml_writer_elem_len(writer, writer->elements,
+                                   writer->nb_elements);
+    if (ret < 0)
+        return ret;
+    ebml_writer_elem_write(writer->elements, pb);
+    return 0;
+}
+
 static void mkv_add_seekhead_entry(MatroskaMuxContext *mkv, uint32_t elementid,
                                    uint64_t filepos)
 {