diff mbox

[FFmpeg-devel] lavf/segment: provide a virtual AVIOContext representing all the segments

Message ID 20160929222623.86968-1-rodger.combs@gmail.com
State Superseded
Headers show

Commit Message

Rodger Combs Sept. 29, 2016, 10:26 p.m. UTC
This allows the use of muxers like matroska, which attempt to seek even
when an AVIOContext doesn't set `seekable`, without concern for a rouge
seek leading the muxer to overwrite the wrong data in a later segment.
---
 doc/muxers.texi       |  17 ++++
 libavformat/segment.c | 276 +++++++++++++++++++++++++++++++++++++-------------
 libavformat/version.h |   2 +-
 3 files changed, 223 insertions(+), 72 deletions(-)
diff mbox

Patch

diff --git a/doc/muxers.texi b/doc/muxers.texi
index 2c937c7..fd2136e 100644
--- a/doc/muxers.texi
+++ b/doc/muxers.texi
@@ -1390,6 +1390,23 @@  argument must be a time duration specification, and defaults to 0.
 If enabled, write an empty segment if there are no packets during the period a
 segment would usually span. Otherwise, the segment will be filled with the next
 packet written. Defaults to @code{0}.
+
+@item individual_header_trailer @var{1|0}
+Write each segment as an individual distinct file in the underlying format.
+When this is set to @code{0}, the segments are treated as a single continuous
+stream. When set to @code{1} (the default), each individual segment receives
+a header and trailer, so they can be played independently.
+
+@item segment_header_filename @var{name}
+Write the global header to a separate file. Requires
+@var{individual_header_trailer} to be @code{0}. If not set, the global header
+will be written to the first file along with actual segment data.
+
+@item segment_seekback @var{1|0}
+Allow the muxer to seek back and overwrite data from previous segments. Requires
+@var{individual_header_trailer} to be @code{0}. If set to @code{0} (the default),
+the underlying muxer will be unable to seek back into previous segments, so they
+can be relied upon not to change once written.
 @end table
 
 @subsection Examples
diff --git a/libavformat/segment.c b/libavformat/segment.c
index 33a5cf0..4858e9c 100644
--- a/libavformat/segment.c
+++ b/libavformat/segment.c
@@ -51,8 +51,10 @@  typedef struct SegmentListEntry {
     int64_t start_pts;
     int64_t offset_pts;
     char *filename;
+    char *full_filename;
     struct SegmentListEntry *next;
     int64_t last_duration;
+    size_t start_offset;
 } SegmentListEntry;
 
 typedef enum {
@@ -125,7 +127,13 @@  typedef struct SegmentContext {
 
     SegmentListEntry cur_entry;
     SegmentListEntry *segment_list_entries;
+    SegmentListEntry *segment_list_entries_all;
     SegmentListEntry *segment_list_entries_end;
+    SegmentListEntry *segment_list_entry_writing;
+    int seekback;          ///< allow seeking back to previous segments
+    AVIOContext *cur_pb;   ///< current segment put-byte context
+    size_t write_offset;
+    size_t max_offset;
 } SegmentContext;
 
 static void print_csv_escaped_str(AVIOContext *ctx, const char *str)
@@ -144,6 +152,123 @@  static void print_csv_escaped_str(AVIOContext *ctx, const char *str)
         avio_w8(ctx, '"');
 }
 
+static int64_t virtual_seek(void *priv, int64_t target, int whence)
+{
+    AVFormatContext *s = priv;
+    SegmentContext *seg = s->priv_data;
+    SegmentListEntry *it, *current = NULL;
+    int64_t offset = target;
+    int64_t ret;
+
+    if (whence != SEEK_SET)
+        return AVERROR(EINVAL);
+    if (offset < 0)
+        return AVERROR(EINVAL);
+
+    if (offset >= seg->max_offset) {
+        ff_format_io_close(s, &seg->cur_pb);
+        seg->write_offset = offset;
+        return offset;
+    }
+
+    if (seg->cur_entry.start_offset <= offset) {
+        current = &seg->cur_entry;
+    } else {
+        for (it = seg->segment_list_entries_all; it; it = it->next) {
+            if (it->start_offset <= offset)
+                current = it;
+            else if (it->start_offset > offset)
+                break;
+        }
+    }
+
+    offset -= current->start_offset;
+
+    if (current != seg->segment_list_entry_writing) {
+        int is_seekback = (current != &seg->cur_entry) && seg->segment_list_entries;
+        char *new_filename;
+        AVIOContext *new_ctx = NULL;
+        AVDictionary *options = NULL;
+
+        if (!seg->seekback && is_seekback)
+            return AVERROR(EINVAL);
+
+        new_filename = current->full_filename;
+
+        if (new_filename) {
+            if (is_seekback)
+                av_dict_set_int(&options, "truncate", 0, 0);
+            if ((ret = s->io_open(s, &new_ctx, new_filename, AVIO_FLAG_WRITE, &options)) < 0) {
+                av_log(s, AV_LOG_ERROR, "Failed to seek into segment '%s'\n", new_filename);
+                return ret;
+            }
+        }
+
+        ff_format_io_close(s, &seg->cur_pb);
+        seg->cur_pb = new_ctx;
+        seg->segment_list_entry_writing = current;
+    }
+
+    if (seg->cur_pb)
+        if ((ret = avio_seek(seg->cur_pb, offset, SEEK_SET)) < 0)
+            return ret;
+
+    seg->write_offset = offset;
+
+    return target;
+}
+
+static int virtual_write(void *priv, uint8_t *buf, int buf_size)
+{
+    AVFormatContext *s = priv;
+    SegmentContext *seg = s->priv_data;
+    int ret = 0;
+    int written = 0;
+
+    while (written < buf_size) {
+        SegmentListEntry *cur = seg->segment_list_entry_writing;
+        size_t start = cur->start_offset + seg->write_offset;
+        SegmentListEntry *next = cur->next ? cur->next : (cur == &seg->cur_entry ? NULL : &seg->cur_entry);
+        size_t end = next ? next->start_offset : SIZE_MAX;
+        int to_write = FFMIN(end - start, buf_size - written);
+        if (seg->cur_pb)
+            avio_write(seg->cur_pb, buf, to_write);
+        buf += to_write;
+        written += to_write;
+        seg->write_offset += to_write;
+        if (written < buf_size)
+            if ((ret = virtual_seek(s, end, SEEK_SET)) < 0)
+                return ret;
+    }
+
+    return written;
+}
+
+static void virtual_close(AVFormatContext *s)
+{
+    SegmentContext *seg = s->priv_data;
+    ff_format_io_close(s, &seg->cur_pb);
+    av_freep(&seg->avf->pb->buffer);
+    av_freep(&seg->avf->pb);
+}
+
+static int open_virtual_ctx(AVFormatContext *s, AVIOContext **ctx)
+{
+    SegmentContext *seg = s->priv_data;
+    int buf_size = 32768;
+    uint8_t *buf = av_malloc(buf_size);
+    if (!buf)
+        return AVERROR(ENOMEM);
+    *ctx = avio_alloc_context(buf, buf_size, AVIO_FLAG_WRITE, s, NULL,
+                              virtual_write, virtual_seek);
+    if (!*ctx) {
+        av_free(buf);
+        return AVERROR(ENOMEM);
+    }
+    (*ctx)->seekable = seg->seekback;
+    return virtual_seek(s, 0, SEEK_SET);
+}
+
 static int segment_mux_init(AVFormatContext *s)
 {
     SegmentContext *seg = s->priv_data;
@@ -193,7 +318,6 @@  static int set_segment_filename(AVFormatContext *s)
     SegmentContext *seg = s->priv_data;
     AVFormatContext *oc = seg->avf;
     size_t size;
-    int ret;
 
     if (seg->segment_idx_wrap)
         seg->segment_idx %= seg->segment_idx_wrap;
@@ -212,13 +336,18 @@  static int set_segment_filename(AVFormatContext *s)
         return AVERROR(EINVAL);
     }
 
+    seg->cur_entry.full_filename = av_strdup(oc->filename);
+    if (!seg->cur_entry.full_filename)
+        return AVERROR(ENOMEM);
+
     /* copy modified name in list entry */
     size = strlen(av_basename(oc->filename)) + 1;
     if (seg->entry_prefix)
         size += strlen(seg->entry_prefix);
 
-    if ((ret = av_reallocp(&seg->cur_entry.filename, size)) < 0)
-        return ret;
+    seg->cur_entry.filename = av_mallocz(size);
+    if (!seg->cur_entry.filename)
+        return AVERROR(ENOMEM);
     snprintf(seg->cur_entry.filename, size, "%s%s",
              seg->entry_prefix ? seg->entry_prefix : "",
              av_basename(oc->filename));
@@ -247,12 +376,16 @@  static int segment_start(AVFormatContext *s, int write_header)
     if ((err = set_segment_filename(s)) < 0)
         return err;
 
-    if ((err = s->io_open(s, &oc->pb, oc->filename, AVIO_FLAG_WRITE, NULL)) < 0) {
-        av_log(s, AV_LOG_ERROR, "Failed to open segment '%s'\n", oc->filename);
-        return err;
+    if (seg->individual_header_trailer) {
+        if ((err = s->io_open(s, &oc->pb, oc->filename, AVIO_FLAG_WRITE, NULL)) < 0) {
+            av_log(s, AV_LOG_ERROR, "Failed to open segment '%s'\n", oc->filename);
+            return err;
+        }
+    } else {
+        seg->cur_entry.start_offset += seg->write_offset;
+        if ((err = virtual_seek(s, seg->cur_entry.start_offset, SEEK_SET)) < 0)
+            return err;
     }
-    if (!seg->individual_header_trailer)
-        oc->pb->seekable = 0;
 
     if (oc->oformat->priv_class && oc->priv_data)
         av_opt_set(oc->priv_data, "mpegts_flags", "+resend_headers", 0);
@@ -344,6 +477,7 @@  static int segment_end(AVFormatContext *s, int write_trailer, int is_last)
 {
     SegmentContext *seg = s->priv_data;
     AVFormatContext *oc = seg->avf;
+    SegmentListEntry *entry;
     int ret = 0;
     AVTimecode tc;
     AVRational rate;
@@ -360,29 +494,27 @@  static int segment_end(AVFormatContext *s, int write_trailer, int is_last)
         av_log(s, AV_LOG_ERROR, "Failure occurred when ending segment '%s'\n",
                oc->filename);
 
-    if (seg->list) {
-        if (seg->list_size || seg->list_type == LIST_TYPE_M3U8) {
-            SegmentListEntry *entry = av_mallocz(sizeof(*entry));
-            if (!entry) {
-                ret = AVERROR(ENOMEM);
-                goto end;
-            }
+    entry = av_mallocz(sizeof(*entry));
+    if (!entry) {
+        ret = AVERROR(ENOMEM);
+        goto end;
+    }
+
+    /* append new element */
+    memcpy(entry, &seg->cur_entry, sizeof(*entry));
+    if (!seg->segment_list_entries)
+        seg->segment_list_entries_all->next = seg->segment_list_entries = entry;
+    else
+        seg->segment_list_entries_end->next = entry;
+    seg->segment_list_entries_end = entry;
 
-            /* append new element */
-            memcpy(entry, &seg->cur_entry, sizeof(*entry));
-            entry->filename = av_strdup(entry->filename);
-            if (!seg->segment_list_entries)
-                seg->segment_list_entries = seg->segment_list_entries_end = entry;
-            else
-                seg->segment_list_entries_end->next = entry;
-            seg->segment_list_entries_end = entry;
+    seg->segment_list_entry_writing = NULL;
 
+    if (seg->list) {
+        if (seg->list_size || seg->list_type == LIST_TYPE_M3U8) {
             /* drop first item */
             if (seg->list_size && seg->segment_count >= seg->list_size) {
-                entry = seg->segment_list_entries;
                 seg->segment_list_entries = seg->segment_list_entries->next;
-                av_freep(&entry->filename);
-                av_freep(&entry);
             }
 
             if ((ret = segment_list_open(s)) < 0)
@@ -428,7 +560,8 @@  static int segment_end(AVFormatContext *s, int write_trailer, int is_last)
     }
 
 end:
-    ff_format_io_close(oc, &oc->pb);
+    if (seg->individual_header_trailer)
+        ff_format_io_close(oc, &oc->pb);
 
     return ret;
 }
@@ -550,26 +683,6 @@  end:
     return ret;
 }
 
-static int open_null_ctx(AVIOContext **ctx)
-{
-    int buf_size = 32768;
-    uint8_t *buf = av_malloc(buf_size);
-    if (!buf)
-        return AVERROR(ENOMEM);
-    *ctx = avio_alloc_context(buf, buf_size, AVIO_FLAG_WRITE, NULL, NULL, NULL, NULL);
-    if (!*ctx) {
-        av_free(buf);
-        return AVERROR(ENOMEM);
-    }
-    return 0;
-}
-
-static void close_null_ctxp(AVIOContext **pb)
-{
-    av_freep(&(*pb)->buffer);
-    av_freep(pb);
-}
-
 static int select_reference_stream(AVFormatContext *s)
 {
     SegmentContext *seg = s->priv_data;
@@ -627,9 +740,12 @@  static int select_reference_stream(AVFormatContext *s)
     return 0;
 }
 
-static void seg_free_context(SegmentContext *seg)
+static void seg_free_context(AVFormatContext *s)
 {
+    SegmentContext *seg = s->priv_data;
     ff_format_io_close(seg->avf, &seg->list_pb);
+    if (!seg->individual_header_trailer)
+        virtual_close(s);
     avformat_free_context(seg->avf);
     seg->avf = NULL;
 }
@@ -642,6 +758,14 @@  static int seg_init(AVFormatContext *s)
     int ret;
     int i;
 
+    seg->max_offset = SIZE_MAX;
+
+    if (seg->write_header_trailer && !seg->header_filename) {
+        seg->cur_entry.start_offset = 0;
+    } else {
+        seg->cur_entry.start_offset = SIZE_MAX;
+    }
+
     seg->segment_count = 0;
     if (!seg->write_header_trailer)
         seg->individual_header_trailer = 0;
@@ -743,17 +867,27 @@  static int seg_init(AVFormatContext *s)
         goto fail;
     oc = seg->avf;
 
-    if (seg->write_header_trailer) {
-        if ((ret = s->io_open(s, &oc->pb,
-                              seg->header_filename ? seg->header_filename : oc->filename,
-                              AVIO_FLAG_WRITE, NULL)) < 0) {
+    seg->segment_list_entries_all = av_mallocz(sizeof(*seg->segment_list_entries_all));
+    if (!seg->segment_list_entries_all) {
+        ret = ENOMEM;
+        goto fail;
+    }
+
+    if (seg->header_filename) {
+        seg->segment_list_entries_all->full_filename = av_strdup(seg->header_filename);
+        if (!seg->segment_list_entries_all->full_filename) {
+            ret = ENOMEM;
+            goto fail;
+        }
+    }
+
+    if (seg->individual_header_trailer) {
+        if ((ret = s->io_open(s, &oc->pb, oc->filename, AVIO_FLAG_WRITE, NULL)) < 0) {
             av_log(s, AV_LOG_ERROR, "Failed to open segment '%s'\n", oc->filename);
             goto fail;
         }
-        if (!seg->individual_header_trailer)
-            oc->pb->seekable = 0;
     } else {
-        if ((ret = open_null_ctx(&oc->pb)) < 0)
+        if ((ret = open_virtual_ctx(s, &oc->pb)) < 0)
             goto fail;
     }
 
@@ -783,22 +917,17 @@  static int seg_init(AVFormatContext *s)
         s->avoid_negative_ts = 1;
 
     if (!seg->write_header_trailer || seg->header_filename) {
-        if (seg->header_filename) {
-            av_write_frame(oc, NULL);
-            ff_format_io_close(oc, &oc->pb);
-        } else {
-            close_null_ctxp(&oc->pb);
-        }
-        if ((ret = oc->io_open(oc, &oc->pb, oc->filename, AVIO_FLAG_WRITE, NULL)) < 0)
+        av_write_frame(oc, NULL);
+        seg->cur_entry.start_offset = seg->write_offset;
+        if ((ret = virtual_seek(s, seg->write_offset, SEEK_SET)) < 0)
             goto fail;
-        if (!seg->individual_header_trailer)
-            oc->pb->seekable = 0;
+        ret = 0;
     }
 
 fail:
     av_dict_free(&options);
     if (ret < 0)
-        seg_free_context(seg);
+        seg_free_context(s);
 
     return ret;
 }
@@ -914,7 +1043,7 @@  fail:
     }
 
     if (ret < 0)
-        seg_free_context(seg);
+        seg_free_context(s);
 
     return ret;
 }
@@ -932,13 +1061,17 @@  static int seg_write_trailer(struct AVFormatContext *s)
     if (!seg->write_header_trailer) {
         if ((ret = segment_end(s, 0, 1)) < 0)
             goto fail;
-        if ((ret = open_null_ctx(&oc->pb)) < 0)
-            goto fail;
+
+        seg->max_offset = seg->cur_entry.start_offset + seg->write_offset;
+        virtual_seek(oc->pb, seg->max_offset, SEEK_SET);
         ret = av_write_trailer(oc);
-        close_null_ctxp(&oc->pb);
     } else {
         ret = segment_end(s, 1, 1);
     }
+
+    if (!seg->individual_header_trailer)
+        virtual_close(s);
+
 fail:
     if (seg->list)
         ff_format_io_close(s, &seg->list_pb);
@@ -947,11 +1080,11 @@  fail:
     av_opt_free(seg);
     av_freep(&seg->times);
     av_freep(&seg->frames);
-    av_freep(&seg->cur_entry.filename);
 
-    cur = seg->segment_list_entries;
+    cur = seg->segment_list_entries_all;
     while (cur) {
         next = cur->next;
+        av_freep(&cur->full_filename);
         av_freep(&cur->filename);
         av_free(cur);
         cur = next;
@@ -1005,6 +1138,7 @@  static const AVOption options[] = {
     { "reset_timestamps", "reset timestamps at the begin of each segment", OFFSET(reset_timestamps), AV_OPT_TYPE_BOOL, {.i64 = 0}, 0, 1, E },
     { "initial_offset", "set initial timestamp offset", OFFSET(initial_offset), AV_OPT_TYPE_DURATION, {.i64 = 0}, -INT64_MAX, INT64_MAX, E },
     { "write_empty_segments", "allow writing empty 'filler' segments", OFFSET(write_empty), AV_OPT_TYPE_BOOL, {.i64 = 0}, 0, 1, E },
+    { "segment_seekback", "allow seeking back to previous segments", OFFSET(seekback), AV_OPT_TYPE_INT, {.i64 = 0}, 0, 1, E },
     { NULL },
 };
 
diff --git a/libavformat/version.h b/libavformat/version.h
index b0b5593..848403e 100644
--- a/libavformat/version.h
+++ b/libavformat/version.h
@@ -33,7 +33,7 @@ 
 // Also please add any ticket numbers that you believe might be affected here
 #define LIBAVFORMAT_VERSION_MAJOR  57
 #define LIBAVFORMAT_VERSION_MINOR  51
-#define LIBAVFORMAT_VERSION_MICRO 100
+#define LIBAVFORMAT_VERSION_MICRO 101
 
 #define LIBAVFORMAT_VERSION_INT AV_VERSION_INT(LIBAVFORMAT_VERSION_MAJOR, \
                                                LIBAVFORMAT_VERSION_MINOR, \