matroskadec: export ordered chapters files as a playlist.

2 files changed, 180 insertions, 2 deletions

diff --git a/libavformat/matroska.h b/libavformat/matroska.h
index 8e747e6a9c..c7c2484abf 100644
--- a/libavformat/matroska.h
+++ b/libavformat/matroska.h
@@ -200,6 +200,7 @@
 #define MATROSKA_ID_CHAPTERFLAGHIDDEN   0x98
 #define MATROSKA_ID_CHAPTERFLAGENABLED  0x4598
 #define MATROSKA_ID_CHAPTERPHYSEQUIV    0x63C3
+#define MATROSKA_ID_CHAPTERSEGMENTUID   0x6E67
 
 typedef enum {
   MATROSKA_TRACK_TYPE_NONE     = 0x0,
diff --git a/libavformat/matroskadec.c b/libavformat/matroskadec.c
index 89df095cd8..bb6b17b727 100644
--- a/libavformat/matroskadec.c
+++ b/libavformat/matroskadec.c
@@ -37,18 +37,22 @@
 #include "isom.h"
 #include "rm.h"
 #include "matroska.h"
+#include "playlist.h"
 #include "libavcodec/mpeg4audio.h"
 #include "libavutil/intfloat_readwrite.h"
 #include "libavutil/intreadwrite.h"
 #include "libavutil/avstring.h"
 #include "libavutil/lzo.h"
 #include "libavutil/dict.h"
+#include "libavutil/mathematics.h"
+#include "libavutil/opt.h"
 #if CONFIG_ZLIB
 #include <zlib.h>
 #endif
 #if CONFIG_BZLIB
 #include <bzlib.h>
 #endif
+#include <libgen.h>
 
 typedef enum {
     EBML_NONE,
@@ -169,6 +173,7 @@ typedef struct {
     char    *title;
 
     AVChapter *chapter;
+    EbmlBin segment_uid;
 } MatroskaChapter;
 
 typedef struct {
@@ -213,6 +218,7 @@ typedef struct {
 } MatroskaLevel;
 
 typedef struct {
+    AVClass *class;
     AVFormatContext *ctx;
 
     /* EBML stuff */
@@ -224,6 +230,7 @@ typedef struct {
     uint64_t time_scale;
     double   duration;
     char    *title;
+    EbmlBin segment_uid;
     EbmlList tracks;
     EbmlList attachments;
     EbmlList chapters;
@@ -247,6 +254,10 @@ typedef struct {
 
     /* File has a CUES element, but we defer parsing until it is needed. */
     int cues_parsing_deferred;
+
+    int ordered_chapters;       /* export ordered chapters as a playlist, set by
+                                   a private option */
+    int ordered_ed_found;
 } MatroskaDemuxContext;
 
 typedef struct {
@@ -281,10 +292,10 @@ static EbmlSyntax matroska_info[] = {
     { MATROSKA_ID_TIMECODESCALE,      EBML_UINT,  0, offsetof(MatroskaDemuxContext,time_scale), {.u=1000000} },
     { MATROSKA_ID_DURATION,           EBML_FLOAT, 0, offsetof(MatroskaDemuxContext,duration) },
     { MATROSKA_ID_TITLE,              EBML_UTF8,  0, offsetof(MatroskaDemuxContext,title) },
+    { MATROSKA_ID_SEGMENTUID,         EBML_BIN,   0, offsetof(MatroskaDemuxContext, segment_uid) },
     { MATROSKA_ID_WRITINGAPP,         EBML_NONE },
     { MATROSKA_ID_MUXINGAPP,          EBML_NONE },
     { MATROSKA_ID_DATEUTC,            EBML_NONE },
-    { MATROSKA_ID_SEGMENTUID,         EBML_NONE },
     { 0 }
 };
 
@@ -390,6 +401,7 @@ static EbmlSyntax matroska_chapter_entry[] = {
     { MATROSKA_ID_CHAPTERTIMEEND,     EBML_UINT, 0, offsetof(MatroskaChapter,end), {.u=AV_NOPTS_VALUE} },
     { MATROSKA_ID_CHAPTERUID,         EBML_UINT, 0, offsetof(MatroskaChapter,uid) },
     { MATROSKA_ID_CHAPTERDISPLAY,     EBML_NEST, 0, 0, {.n=matroska_chapter_display} },
+    { MATROSKA_ID_CHAPTERSEGMENTUID,  EBML_BIN,  0, offsetof(MatroskaChapter, segment_uid) },
     { MATROSKA_ID_CHAPTERFLAGHIDDEN,  EBML_NONE },
     { MATROSKA_ID_CHAPTERFLAGENABLED, EBML_NONE },
     { MATROSKA_ID_CHAPTERPHYSEQUIV,   EBML_NONE },
@@ -399,10 +411,10 @@ static EbmlSyntax matroska_chapter_entry[] = {
 
 static EbmlSyntax matroska_chapter[] = {
     { MATROSKA_ID_CHAPTERATOM,        EBML_NEST, sizeof(MatroskaChapter), offsetof(MatroskaDemuxContext,chapters), {.n=matroska_chapter_entry} },
+    { MATROSKA_ID_EDITIONFLAGORDERED, EBML_UINT, 0, offsetof(MatroskaDemuxContext, ordered_ed_found), {.u=0} },
     { MATROSKA_ID_EDITIONUID,         EBML_NONE },
     { MATROSKA_ID_EDITIONFLAGHIDDEN,  EBML_NONE },
     { MATROSKA_ID_EDITIONFLAGDEFAULT, EBML_NONE },
-    { MATROSKA_ID_EDITIONFLAGORDERED, EBML_NONE },
     { 0 }
 };
 
@@ -1269,6 +1281,154 @@ static int matroska_aac_sri(int samplerate)
     return sri;
 }
 
+/* stripped down syntax for segment UID probing */
+static EbmlSyntax matroska_segment_simple[] = {
+    { MATROSKA_ID_INFO,        EBML_NEST, 0, 0, { .n = matroska_info } },
+    { MATROSKA_ID_SEEKHEAD,    EBML_NONE },
+    { MATROSKA_ID_TRACKS,      EBML_NONE },
+    { MATROSKA_ID_ATTACHMENTS, EBML_NONE },
+    { MATROSKA_ID_CHAPTERS,    EBML_NONE },
+    { MATROSKA_ID_CUES,        EBML_NONE },
+    { MATROSKA_ID_TAGS,        EBML_NONE },
+    { MATROSKA_ID_CLUSTER,     EBML_STOP },
+    { 0 }
+};
+
+static EbmlSyntax matroska_segments_simple[] = {
+    { MATROSKA_ID_SEGMENT,        EBML_NEST, 0, 0, { .n = matroska_segment_simple } },
+    { 0 }
+};
+
+static int matroska_setup_ordered_chapters(MatroskaDemuxContext *matroska)
+{
+    struct segment {
+        uint8_t uid[16];
+        char *filename;
+    } *segments = NULL;
+    int i, ret = 0, nb_segments = 0;
+    MatroskaChapter *chapters = matroska->chapters.elem;
+    AVFormatContext *s = matroska->ctx;
+
+    if (matroska->segment_uid.size != 16) {
+        av_log(s, AV_LOG_ERROR, "Invalid or missing segment UID.\n");
+        return AVERROR_INVALIDDATA;
+    }
+
+    /* find out which external segments are needed */
+    for (i = 0; i < matroska->chapters.nb_elem; i++) {
+        MatroskaChapter *c = &chapters[i];
+
+        if (!c->segment_uid.size)
+            continue;
+        if (c->segment_uid.size != 16) {
+            av_log(s, AV_LOG_ERROR, "Invalid segment UID size %d in chapter %d.\n",
+                   c->segment_uid.size, i);
+            goto error;
+        }
+
+        if (memcmp(c->segment_uid.data, matroska->segment_uid.data, c->segment_uid.size)) {
+            struct segment *tmp = av_realloc(segments, (nb_segments + 1)*sizeof(*segments));
+            if (!tmp) {
+                ret = AVERROR(ENOMEM);
+                goto error;
+            }
+            segments = tmp;
+
+            memcpy(segments[nb_segments].uid, c->segment_uid.data, c->segment_uid.size);
+            segments[nb_segments].filename = NULL;
+            nb_segments++;
+        }
+    }
+
+    /* scan all mkv files in the same dir as this -- de-facto standard behavior */
+    if (nb_segments) {
+        URLContext *uc;
+        char filename[1024], *dir;
+        int needed_segments = nb_segments;
+
+        if (!s->filename) {
+            av_log(s, AV_LOG_ERROR, "No filename provided for this file. "
+                                    "Cannot scan for linked segments.\n");
+            ret = AVERROR(EINVAL);
+            goto error;
+        }
+
+        av_strlcpy(filename, s->filename, sizeof(filename));
+        dir = dirname(filename);
+
+        av_log(s, AV_LOG_DEBUG, "Looking for linked segments in %s.\n", dir);
+        ret = ffurl_open(&uc, dir, AVIO_FLAG_DIR | AVIO_FLAG_READ);
+        if (ret < 0)
+            goto error;
+
+        while (needed_segments && ffurl_read(uc, filename, sizeof(filename)) >= 0) {
+            AVFormatContext *s1;
+            MatroskaDemuxContext *m1;
+            AVIOContext *pb;
+            Ebml ebml = { 0 };
+
+            if(!strcmp(s->filename, filename) || !av_match_ext(filename, "mkv"))
+                continue;
+
+            av_log(s, AV_LOG_DEBUG, "Trying file %s.\n", filename);
+            if (avio_open(&pb, filename, AVIO_FLAG_READ) < 0)
+                continue;
+
+            m1 = av_mallocz(sizeof(*m1));
+            m1->ctx = s1 = avformat_alloc_context();
+            s1->pb = pb;
+            if (ebml_parse(m1, ebml_syntax, &ebml) >= 0) {
+                ebml_free(ebml_syntax, &ebml);
+                ebml_parse(m1, matroska_segments_simple, m1);
+                if (m1->segment_uid.size == 16)
+                    for (i = 0; i < nb_segments; i++)
+                        if (!segments[i].filename &&
+                                !memcmp(m1->segment_uid.data, segments[i].uid, m1->segment_uid.size)) {
+                            segments[i].filename = av_strdup(filename);
+                            needed_segments--;
+                            break;
+                        }
+                ebml_free(matroska_segments_simple, m1);
+            }
+            avio_close(pb);
+            av_freep(&m1);
+            av_freep(&s1);
+        }
+        ffurl_close(uc);
+    }
+
+    for (i = 0; i < matroska->chapters.nb_elem; i++) {
+        MatroskaChapter *c = &chapters[i];
+        AVPlaylistElement *elem = NULL;
+        int j;
+
+        if (c->segment_uid.size) {
+            for (j = 0; j < nb_segments; j++)
+                if (!memcmp(segments[j].uid, c->segment_uid.data, c->segment_uid.size)) {
+                    if (segments[j].filename)
+                        elem = ff_playlist_add_entry(s, segments[j].filename);
+                    else
+                        av_log(s, AV_LOG_WARNING, "Linked segment for chapter %d not found.\n", i);
+                    break;
+                }
+        } else
+            elem = ff_playlist_add_entry(matroska->ctx, matroska->ctx->filename);
+
+        if (elem) {
+            if (c->start != AV_NOPTS_VALUE)
+                elem->start = av_rescale_q(c->start, (AVRational){1,1000000000}, AV_TIME_BASE_Q);
+            if (c->end != AV_NOPTS_VALUE)
+                elem->end   = av_rescale_q(c->end, (AVRational){1,1000000000}, AV_TIME_BASE_Q);
+        }
+    }
+
+error:
+    for (i = 0; i < nb_segments; i++)
+        av_freep(&segments[i].filename);
+    av_freep(&segments);
+    return ret;
+}
+
 static int matroska_read_header(AVFormatContext *s, AVFormatParameters *ap)
 {
     MatroskaDemuxContext *matroska = s->priv_data;
@@ -1315,6 +1475,9 @@ static int matroska_read_header(AVFormatContext *s, AVFormatParameters *ap)
                                   * 1000 / AV_TIME_BASE;
     av_dict_set(&s->metadata, "title", matroska->title, 0);
 
+    if (matroska->ordered_chapters && matroska->ordered_ed_found)
+        return matroska_setup_ordered_chapters(matroska);
+
     tracks = matroska->tracks.elem;
     for (i=0; i < matroska->tracks.nb_elem; i++) {
         MatroskaTrack *track = &tracks[i];
@@ -2022,6 +2185,19 @@ static int matroska_read_close(AVFormatContext *s)
     return 0;
 }
 
+#define OFFSET(x) offsetof(MatroskaDemuxContext, x)
+static const AVOption options[] = {
+    { "ordered_chapters", "Export files with ordered chapters as playlists.", OFFSET(ordered_chapters), FF_OPT_TYPE_INT, { 0 }, 0, 1, AV_OPT_FLAG_DECODING_PARAM },
+    { NULL },
+};
+
+static const AVClass matroska_demuxer_class = {
+    .class_name     = "Matroska demuxer",
+    .item_name      = av_default_item_name,
+    .option         = options,
+    .version        = LIBAVUTIL_VERSION_INT,
+};
+
 AVInputFormat ff_matroska_demuxer = {
     .name           = "matroska,webm",
     .long_name      = NULL_IF_CONFIG_SMALL("Matroska/WebM file format"),
@@ -2031,4 +2207,5 @@ AVInputFormat ff_matroska_demuxer = {
     .read_packet    = matroska_read_packet,
     .read_close     = matroska_read_close,
     .read_seek      = matroska_read_seek,
+    .priv_class     = &matroska_demuxer_class,
 };