segment: fix NULL pointer dereference in seg_write_header()
[ffmpeg.git] / libavformat / segment.c
1 /*
2  * Generic segmenter
3  * Copyright (c) 2011, Luca Barbato
4  *
5  * This file is part of Libav.
6  *
7  * Libav is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU Lesser General Public
9  * License as published by the Free Software Foundation; either
10  * version 2.1 of the License, or (at your option) any later version.
11  *
12  * Libav is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15  * Lesser General Public License for more details.
16  *
17  * You should have received a copy of the GNU Lesser General Public
18  * License along with Libav; if not, write to the Free Software
19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
20  */
21
22 #include <float.h>
23
24 #include "avformat.h"
25 #include "internal.h"
26
27 #include "libavutil/log.h"
28 #include "libavutil/opt.h"
29 #include "libavutil/avstring.h"
30 #include "libavutil/parseutils.h"
31 #include "libavutil/mathematics.h"
32
33 typedef struct {
34     const AVClass *class;  /**< Class for private options. */
35     int number;
36     AVOutputFormat *oformat;
37     AVFormatContext *avf;
38     char *format;          /**< Set by a private option. */
39     char *list;            /**< Set by a private option. */
40     int  list_type;        /**< Set by a private option. */
41     float time;            /**< Set by a private option. */
42     int  size;             /**< Set by a private option. */
43     int  wrap;             /**< Set by a private option. */
44     int  individual_header_trailer; /**< Set by a private option. */
45     int  write_header_trailer; /**< Set by a private option. */
46     int64_t offset_time;
47     int64_t recording_time;
48     int has_video;
49     AVIOContext *pb;
50 } SegmentContext;
51
52 enum {
53     LIST_FLAT,
54     LIST_HLS
55 };
56
57 static int segment_mux_init(AVFormatContext *s)
58 {
59     SegmentContext *seg = s->priv_data;
60     AVFormatContext *oc;
61     int i;
62
63     seg->avf = oc = avformat_alloc_context();
64     if (!oc)
65         return AVERROR(ENOMEM);
66
67     oc->oformat            = seg->oformat;
68     oc->interrupt_callback = s->interrupt_callback;
69
70     for (i = 0; i < s->nb_streams; i++) {
71         AVStream *st;
72         if (!(st = avformat_new_stream(oc, NULL)))
73             return AVERROR(ENOMEM);
74         avcodec_copy_context(st->codec, s->streams[i]->codec);
75         st->sample_aspect_ratio = s->streams[i]->sample_aspect_ratio;
76     }
77
78     return 0;
79 }
80
81 static int segment_hls_window(AVFormatContext *s, int last)
82 {
83     SegmentContext *seg = s->priv_data;
84     int i, ret = 0;
85     char buf[1024];
86
87     if ((ret = avio_open2(&seg->pb, seg->list, AVIO_FLAG_WRITE,
88                               &s->interrupt_callback, NULL)) < 0)
89         goto fail;
90
91     avio_printf(seg->pb, "#EXTM3U\n");
92     avio_printf(seg->pb, "#EXT-X-VERSION:3\n");
93     avio_printf(seg->pb, "#EXT-X-TARGETDURATION:%d\n", (int)seg->time);
94     avio_printf(seg->pb, "#EXT-X-MEDIA-SEQUENCE:%d\n",
95                 FFMAX(0, seg->number - seg->size));
96
97     for (i = FFMAX(0, seg->number - seg->size);
98          i < seg->number; i++) {
99         avio_printf(seg->pb, "#EXTINF:%d,\n", (int)seg->time);
100         av_get_frame_filename(buf, sizeof(buf), s->filename, i);
101         avio_printf(seg->pb, "%s\n", buf);
102     }
103
104     if (last)
105         avio_printf(seg->pb, "#EXT-X-ENDLIST\n");
106 fail:
107     avio_closep(&seg->pb);
108     return ret;
109 }
110
111 static int segment_start(AVFormatContext *s, int write_header)
112 {
113     SegmentContext *c = s->priv_data;
114     AVFormatContext *oc = c->avf;
115     int err = 0;
116
117     if (write_header) {
118         avformat_free_context(oc);
119         c->avf = NULL;
120         if ((err = segment_mux_init(s)) < 0)
121             return err;
122         oc = c->avf;
123     }
124
125     if (c->wrap)
126         c->number %= c->wrap;
127
128     if (av_get_frame_filename(oc->filename, sizeof(oc->filename),
129                               s->filename, c->number++) < 0)
130         return AVERROR(EINVAL);
131
132     if ((err = avio_open2(&oc->pb, oc->filename, AVIO_FLAG_WRITE,
133                           &s->interrupt_callback, NULL)) < 0)
134         return err;
135
136     if (oc->oformat->priv_class && oc->priv_data)
137         av_opt_set(oc->priv_data, "resend_headers", "1", 0); /* mpegts specific */
138
139     if (write_header) {
140         if ((err = avformat_write_header(oc, NULL)) < 0)
141             return err;
142     }
143
144     return 0;
145 }
146
147 static int segment_end(AVFormatContext *oc, int write_trailer)
148 {
149     int ret = 0;
150
151     av_write_frame(oc, NULL); /* Flush any buffered data (fragmented mp4) */
152     if (write_trailer)
153         av_write_trailer(oc);
154     avio_close(oc->pb);
155
156     return ret;
157 }
158
159 static int open_null_ctx(AVIOContext **ctx)
160 {
161     int buf_size = 32768;
162     uint8_t *buf = av_malloc(buf_size);
163     if (!buf)
164         return AVERROR(ENOMEM);
165     *ctx = avio_alloc_context(buf, buf_size, AVIO_FLAG_WRITE, NULL, NULL, NULL, NULL);
166     if (!*ctx) {
167         av_free(buf);
168         return AVERROR(ENOMEM);
169     }
170     return 0;
171 }
172
173 static void close_null_ctx(AVIOContext *pb)
174 {
175     av_free(pb->buffer);
176     av_free(pb);
177 }
178
179 static int seg_write_header(AVFormatContext *s)
180 {
181     SegmentContext *seg = s->priv_data;
182     AVFormatContext *oc = NULL;
183     int ret, i;
184
185     seg->number = 0;
186     seg->offset_time = 0;
187     seg->recording_time = seg->time * 1000000;
188     if (!seg->write_header_trailer)
189         seg->individual_header_trailer = 0;
190
191     if (seg->list && seg->list_type != LIST_HLS)
192         if ((ret = avio_open2(&seg->pb, seg->list, AVIO_FLAG_WRITE,
193                               &s->interrupt_callback, NULL)) < 0)
194             goto fail;
195
196     for (i = 0; i < s->nb_streams; i++)
197         seg->has_video +=
198             (s->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO);
199
200     if (seg->has_video > 1)
201         av_log(s, AV_LOG_WARNING,
202                "More than a single video stream present, "
203                "expect issues decoding it.\n");
204
205     seg->oformat = av_guess_format(seg->format, s->filename, NULL);
206
207     if (!seg->oformat) {
208         ret = AVERROR_MUXER_NOT_FOUND;
209         goto fail;
210     }
211     if (seg->oformat->flags & AVFMT_NOFILE) {
212         av_log(s, AV_LOG_ERROR, "format %s not supported.\n",
213                seg->oformat->name);
214         ret = AVERROR(EINVAL);
215         goto fail;
216     }
217
218     if ((ret = segment_mux_init(s)) < 0)
219         goto fail;
220     oc = seg->avf;
221
222     if (av_get_frame_filename(oc->filename, sizeof(oc->filename),
223                               s->filename, seg->number++) < 0) {
224         ret = AVERROR(EINVAL);
225         goto fail;
226     }
227
228     if (seg->write_header_trailer) {
229         if ((ret = avio_open2(&oc->pb, oc->filename, AVIO_FLAG_WRITE,
230                               &s->interrupt_callback, NULL)) < 0)
231             goto fail;
232     } else {
233         if ((ret = open_null_ctx(&oc->pb)) < 0)
234             goto fail;
235     }
236
237     if ((ret = avformat_write_header(oc, NULL)) < 0) {
238         avio_close(oc->pb);
239         goto fail;
240     }
241
242     if (!seg->write_header_trailer) {
243         close_null_ctx(oc->pb);
244         if ((ret = avio_open2(&oc->pb, oc->filename, AVIO_FLAG_WRITE,
245                               &s->interrupt_callback, NULL)) < 0)
246             goto fail;
247     }
248
249     if (seg->list) {
250         if (seg->list_type == LIST_HLS) {
251             if ((ret = segment_hls_window(s, 0)) < 0)
252                 goto fail;
253         } else {
254             avio_printf(seg->pb, "%s\n", oc->filename);
255             avio_flush(seg->pb);
256         }
257     }
258
259 fail:
260     if (ret) {
261         if (seg->list)
262             avio_close(seg->pb);
263         if (seg->avf)
264             avformat_free_context(seg->avf);
265     }
266     return ret;
267 }
268
269 static int seg_write_packet(AVFormatContext *s, AVPacket *pkt)
270 {
271     SegmentContext *seg = s->priv_data;
272     AVFormatContext *oc = seg->avf;
273     AVStream *st = s->streams[pkt->stream_index];
274     int64_t end_pts = seg->recording_time * seg->number;
275     int ret;
276
277     if ((seg->has_video && st->codec->codec_type == AVMEDIA_TYPE_VIDEO) &&
278         av_compare_ts(pkt->pts, st->time_base,
279                       end_pts, AV_TIME_BASE_Q) >= 0 &&
280         pkt->flags & AV_PKT_FLAG_KEY) {
281
282         av_log(s, AV_LOG_DEBUG, "Next segment starts at %d %"PRId64"\n",
283                pkt->stream_index, pkt->pts);
284
285         ret = segment_end(oc, seg->individual_header_trailer);
286
287         if (!ret)
288             ret = segment_start(s, seg->individual_header_trailer);
289
290         if (ret)
291             goto fail;
292
293         oc = seg->avf;
294
295         if (seg->list) {
296             if (seg->list_type == LIST_HLS) {
297                 if ((ret = segment_hls_window(s, 0)) < 0)
298                     goto fail;
299             } else {
300                 avio_printf(seg->pb, "%s\n", oc->filename);
301                 avio_flush(seg->pb);
302                 if (seg->size && !(seg->number % seg->size)) {
303                     avio_closep(&seg->pb);
304                     if ((ret = avio_open2(&seg->pb, seg->list, AVIO_FLAG_WRITE,
305                                           &s->interrupt_callback, NULL)) < 0)
306                         goto fail;
307                 }
308             }
309         }
310     }
311
312     ret = ff_write_chained(oc, pkt->stream_index, pkt, s);
313
314 fail:
315     if (ret < 0) {
316         if (seg->list)
317             avio_close(seg->pb);
318         avformat_free_context(oc);
319     }
320
321     return ret;
322 }
323
324 static int seg_write_trailer(struct AVFormatContext *s)
325 {
326     SegmentContext *seg = s->priv_data;
327     AVFormatContext *oc = seg->avf;
328     int ret;
329     if (!seg->write_header_trailer) {
330         if ((ret = segment_end(oc, 0)) < 0)
331             goto fail;
332         open_null_ctx(&oc->pb);
333         ret = av_write_trailer(oc);
334         close_null_ctx(oc->pb);
335     } else {
336         ret = segment_end(oc, 1);
337     }
338
339     if (ret < 0)
340         goto fail;
341
342     if (seg->list && seg->list_type == LIST_HLS) {
343         if ((ret = segment_hls_window(s, 1) < 0))
344             goto fail;
345     }
346
347 fail:
348     avio_close(seg->pb);
349     avformat_free_context(oc);
350     return ret;
351 }
352
353 #define OFFSET(x) offsetof(SegmentContext, x)
354 #define E AV_OPT_FLAG_ENCODING_PARAM
355 static const AVOption options[] = {
356     { "segment_format",    "container format used for the segments",  OFFSET(format),  AV_OPT_TYPE_STRING, {.str = NULL},  0, 0,       E },
357     { "segment_time",      "segment length in seconds",               OFFSET(time),    AV_OPT_TYPE_FLOAT,  {.dbl = 2},     0, FLT_MAX, E },
358     { "segment_list",      "output the segment list",                 OFFSET(list),    AV_OPT_TYPE_STRING, {.str = NULL},  0, 0,       E },
359     { "segment_list_size", "maximum number of playlist entries",      OFFSET(size),    AV_OPT_TYPE_INT,    {.i64 = 5},     0, INT_MAX, E },
360     { "segment_list_type", "segment list format",                     OFFSET(list_type),    AV_OPT_TYPE_INT,    {.i64 = LIST_FLAT},     0, 2, E, "list_type" },
361     {   "flat",            "plain list (default)",                    0,               AV_OPT_TYPE_CONST,  {.i64 = LIST_FLAT}, 0, 0, E, "list_type" },
362     {   "hls",             "Apple HTTP Live Streaming compatible",    0,               AV_OPT_TYPE_CONST,  {.i64 = LIST_HLS},  0, 0, E, "list_type" },
363     { "segment_wrap",      "number after which the index wraps",      OFFSET(wrap),    AV_OPT_TYPE_INT,    {.i64 = 0},     0, INT_MAX, E },
364     { "individual_header_trailer", "write header/trailer to each segment", OFFSET(individual_header_trailer), AV_OPT_TYPE_INT, {.i64 = 1}, 0, 1, E },
365     { "write_header_trailer", "write a header to the first segment and a trailer to the last one", OFFSET(write_header_trailer), AV_OPT_TYPE_INT, {.i64 = 1}, 0, 1, E },
366     { NULL },
367 };
368
369 static const AVClass seg_class = {
370     .class_name = "segment muxer",
371     .item_name  = av_default_item_name,
372     .option     = options,
373     .version    = LIBAVUTIL_VERSION_INT,
374 };
375
376
377 AVOutputFormat ff_segment_muxer = {
378     .name           = "segment",
379     .long_name      = NULL_IF_CONFIG_SMALL("segment"),
380     .priv_data_size = sizeof(SegmentContext),
381     .flags          = AVFMT_NOFILE,
382     .write_header   = seg_write_header,
383     .write_packet   = seg_write_packet,
384     .write_trailer  = seg_write_trailer,
385     .priv_class     = &seg_class,
386 };