segment: Free and reinit the muxer before calling avformat_write_header

This makes sure the muxers are set up in the way they expect
with no data left around from the previous run (which could
cause various issues including memory leaks, depending on the chaine
muxer).

This fixes memory leaks with the mpegts and flv muxers. It also
makes the usage of chained muxers correct.

Signed-off-by: Martin Storsjö <martin@martin.st>

1 files changed, 42 insertions, 22 deletions

diff --git a/libavformat/segment.c b/libavformat/segment.c
index 20adb73d24..efd8f8246d 100644
--- a/libavformat/segment.c
+++ b/libavformat/segment.c
@@ -33,6 +33,7 @@
 typedef struct {
     const AVClass *class;  /**< Class for private options. */
     int number;
+    AVOutputFormat *oformat;
     AVFormatContext *avf;
     char *format;          /**< Set by a private option. */
     char *list;            /**< Set by a private option. */
@@ -45,12 +46,42 @@ typedef struct {
     AVIOContext *pb;
 } SegmentContext;
 
+static int segment_mux_init(AVFormatContext *s)
+{
+    SegmentContext *seg = s->priv_data;
+    AVFormatContext *oc;
+    int i;
+
+    seg->avf = oc = avformat_alloc_context();
+    if (!oc)
+        return AVERROR(ENOMEM);
+
+    oc->oformat            = seg->oformat;
+    oc->interrupt_callback = s->interrupt_callback;
+
+    for (i = 0; i < s->nb_streams; i++) {
+        AVStream *st;
+        if (!(st = avformat_new_stream(oc, NULL)))
+            return AVERROR(ENOMEM);
+        avcodec_copy_context(st->codec, s->streams[i]->codec);
+        st->sample_aspect_ratio = s->streams[i]->sample_aspect_ratio;
+    }
+
+    return 0;
+}
+
 static int segment_start(AVFormatContext *s)
 {
     SegmentContext *c = s->priv_data;
     AVFormatContext *oc = c->avf;
     int err = 0;
 
+    avformat_free_context(oc);
+    c->avf = NULL;
+    if ((err = segment_mux_init(s)) < 0)
+        return err;
+    oc = c->avf;
+
     if (c->wrap)
         c->number %= c->wrap;
 
@@ -81,18 +112,13 @@ static int segment_end(AVFormatContext *oc)
 static int seg_write_header(AVFormatContext *s)
 {
     SegmentContext *seg = s->priv_data;
-    AVFormatContext *oc;
+    AVFormatContext *oc = NULL;
     int ret, i;
 
     seg->number = 0;
     seg->offset_time = 0;
     seg->recording_time = seg->time * 1000000;
 
-    oc = avformat_alloc_context();
-
-    if (!oc)
-        return AVERROR(ENOMEM);
-
     if (seg->list)
         if ((ret = avio_open2(&seg->pb, seg->list, AVIO_FLAG_WRITE,
                               &s->interrupt_callback, NULL)) < 0)
@@ -107,31 +133,22 @@ static int seg_write_header(AVFormatContext *s)
                "More than a single video stream present, "
                "expect issues decoding it.\n");
 
-    oc->oformat = av_guess_format(seg->format, s->filename, NULL);
+    seg->oformat = av_guess_format(seg->format, s->filename, NULL);
 
-    if (!oc->oformat) {
+    if (!seg->oformat) {
         ret = AVERROR_MUXER_NOT_FOUND;
         goto fail;
     }
-    if (oc->oformat->flags & AVFMT_NOFILE) {
+    if (seg->oformat->flags & AVFMT_NOFILE) {
         av_log(s, AV_LOG_ERROR, "format %s not supported.\n",
                oc->oformat->name);
         ret = AVERROR(EINVAL);
         goto fail;
     }
 
-    oc->interrupt_callback = s->interrupt_callback;
-    seg->avf = oc;
-
-    for (i = 0; i < s->nb_streams; i++) {
-        AVStream *st;
-        if (!(st = avformat_new_stream(oc, NULL))) {
-            ret = AVERROR(ENOMEM);
-            goto fail;
-        }
-        avcodec_copy_context(st->codec, s->streams[i]->codec);
-        st->sample_aspect_ratio = s->streams[i]->sample_aspect_ratio;
-    }
+    if ((ret = segment_mux_init(s)) < 0)
+        goto fail;
+    oc = seg->avf;
 
     if (av_get_frame_filename(oc->filename, sizeof(oc->filename),
                               s->filename, seg->number++) < 0) {
@@ -157,7 +174,8 @@ fail:
     if (ret) {
         if (seg->list)
             avio_close(seg->pb);
-        avformat_free_context(oc);
+        if (seg->avf)
+            avformat_free_context(seg->avf);
     }
     return ret;
 }
@@ -186,6 +204,8 @@ static int seg_write_packet(AVFormatContext *s, AVPacket *pkt)
         if (ret)
             goto fail;
 
+        oc = seg->avf;
+
         if (seg->list) {
             avio_printf(seg->pb, "%s\n", oc->filename);
             avio_flush(seg->pb);