1/*
2 * Copyright (c) 2015, Vignesh Venkatasubramanian
3 *
4 * This file is part of FFmpeg.
5 *
6 * FFmpeg is free software; you can redistribute it and/or
7 * modify it under the terms of the GNU Lesser General Public
8 * License as published by the Free Software Foundation; either
9 * version 2.1 of the License, or (at your option) any later version.
10 *
11 * FFmpeg is distributed in the hope that it will be useful,
12 * but WITHOUT ANY WARRANTY; without even the implied warranty of
13 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14 * Lesser General Public License for more details.
15 *
16 * You should have received a copy of the GNU Lesser General Public
17 * License along with FFmpeg; if not, write to the Free Software
18 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19 */
20
21/**
22 * @file WebM Chunk Muxer
23 * The chunk muxer enables writing WebM Live chunks where there is a header
24 * chunk, followed by data chunks where each Cluster is written out as a Chunk.
25 */
26
27#include "avformat.h"
28#include "avio.h"
29#include "avio_internal.h"
30#include "internal.h"
31#include "mux.h"
32
33#include "libavutil/log.h"
34#include "libavutil/opt.h"
35#include "libavutil/mathematics.h"
36
37#define MAX_FILENAME_SIZE 1024
38
39typedef struct WebMChunkContext {
40    const AVClass *class;
41    char *header_filename;
42    int chunk_duration;
43    int chunk_index;
44    char *http_method;
45    uint64_t duration_written;
46    int64_t prev_pts;
47    AVFormatContext *avf;
48    int header_written;
49} WebMChunkContext;
50
51static int webm_chunk_init(AVFormatContext *s)
52{
53    WebMChunkContext *wc = s->priv_data;
54    const AVOutputFormat *oformat;
55    AVFormatContext *oc;
56    AVStream *st, *ost = s->streams[0];
57    AVDictionary *dict = NULL;
58    int ret;
59
60    // DASH Streams can only have one track per file.
61    if (s->nb_streams != 1)
62        return AVERROR(EINVAL);
63
64    if (!wc->header_filename) {
65        av_log(s, AV_LOG_ERROR, "No header filename provided\n");
66        return AVERROR(EINVAL);
67    }
68
69    wc->prev_pts = AV_NOPTS_VALUE;
70
71    oformat = av_guess_format("webm", s->url, "video/webm");
72    if (!oformat)
73        return AVERROR_MUXER_NOT_FOUND;
74
75    ret = avformat_alloc_output_context2(&wc->avf, oformat, NULL, NULL);
76    if (ret < 0)
77        return ret;
78    oc = wc->avf;
79
80    ff_format_set_url(oc, wc->header_filename);
81    wc->header_filename = NULL;
82
83    oc->interrupt_callback    = s->interrupt_callback;
84    oc->max_delay             = s->max_delay;
85    oc->flags                 = s->flags & ~AVFMT_FLAG_FLUSH_PACKETS;
86    oc->strict_std_compliance = s->strict_std_compliance;
87    oc->avoid_negative_ts     = s->avoid_negative_ts;
88
89    oc->flush_packets         = 0;
90
91    if ((ret = av_dict_copy(&oc->metadata, s->metadata, 0)) < 0)
92        return ret;
93
94    if (!(st = avformat_new_stream(oc, NULL)))
95        return AVERROR(ENOMEM);
96
97    if ((ret = ff_stream_encode_params_copy(st, ost)) < 0)
98        return ret;
99
100    if (wc->http_method)
101        if ((ret = av_dict_set(&dict, "method", wc->http_method, 0)) < 0)
102            return ret;
103    ret = s->io_open(s, &oc->pb, oc->url, AVIO_FLAG_WRITE, &dict);
104    av_dict_free(&dict);
105    if (ret < 0)
106        return ret;
107    oc->pb->seekable = 0;
108
109    if ((ret = av_dict_set_int(&dict, "dash", 1, 0))   < 0 ||
110        (ret = av_dict_set_int(&dict, "cluster_time_limit",
111                               wc->chunk_duration, 0)) < 0 ||
112        (ret = av_dict_set_int(&dict, "live", 1, 0))   < 0)
113        goto fail;
114
115    ret = avformat_init_output(oc, &dict);
116fail:
117    av_dict_free(&dict);
118    if (ret < 0)
119        return ret;
120
121    // Copy the timing info back to the original stream
122    // so that the timestamps of the packets are directly usable
123    avpriv_set_pts_info(ost, st->pts_wrap_bits, st->time_base.num,
124                                                st->time_base.den);
125
126    // This ensures that the timestamps will already be properly shifted
127    // when the packets arrive here, so we don't need to shift again.
128    s->avoid_negative_ts  = oc->avoid_negative_ts;
129    ffformatcontext(s)->avoid_negative_ts_use_pts =
130        ffformatcontext(oc)->avoid_negative_ts_use_pts;
131    oc->avoid_negative_ts = AVFMT_AVOID_NEG_TS_DISABLED;
132
133    return 0;
134}
135
136static int get_chunk_filename(AVFormatContext *s, char filename[MAX_FILENAME_SIZE])
137{
138    WebMChunkContext *wc = s->priv_data;
139    if (!filename) {
140        return AVERROR(EINVAL);
141    }
142    if (av_get_frame_filename(filename, MAX_FILENAME_SIZE,
143                              s->url, wc->chunk_index - 1) < 0) {
144        av_log(s, AV_LOG_ERROR, "Invalid chunk filename template '%s'\n", s->url);
145        return AVERROR(EINVAL);
146    }
147    return 0;
148}
149
150static int webm_chunk_write_header(AVFormatContext *s)
151{
152    WebMChunkContext *wc = s->priv_data;
153    AVFormatContext *oc = wc->avf;
154    int ret;
155
156    ret = avformat_write_header(oc, NULL);
157    ff_format_io_close(s, &oc->pb);
158    wc->header_written = 1;
159    if (ret < 0)
160        return ret;
161    return 0;
162}
163
164static int chunk_start(AVFormatContext *s)
165{
166    WebMChunkContext *wc = s->priv_data;
167    AVFormatContext *oc = wc->avf;
168    int ret;
169
170    ret = avio_open_dyn_buf(&oc->pb);
171    if (ret < 0)
172        return ret;
173    wc->chunk_index++;
174    return 0;
175}
176
177static int chunk_end(AVFormatContext *s, int flush)
178{
179    WebMChunkContext *wc = s->priv_data;
180    AVFormatContext *oc = wc->avf;
181    int ret;
182    int buffer_size;
183    uint8_t *buffer;
184    AVIOContext *pb;
185    char filename[MAX_FILENAME_SIZE];
186    AVDictionary *options = NULL;
187
188    if (!oc->pb)
189        return 0;
190
191    if (flush)
192        // Flush the cluster in WebM muxer.
193        av_write_frame(oc, NULL);
194    buffer_size = avio_close_dyn_buf(oc->pb, &buffer);
195    oc->pb = NULL;
196    ret = get_chunk_filename(s, filename);
197    if (ret < 0)
198        goto fail;
199    if (wc->http_method)
200        if ((ret = av_dict_set(&options, "method", wc->http_method, 0)) < 0)
201            goto fail;
202    ret = s->io_open(s, &pb, filename, AVIO_FLAG_WRITE, &options);
203    av_dict_free(&options);
204    if (ret < 0)
205        goto fail;
206    avio_write(pb, buffer, buffer_size);
207    ff_format_io_close(s, &pb);
208fail:
209    av_free(buffer);
210    return (ret < 0) ? ret : 0;
211}
212
213static int webm_chunk_write_packet(AVFormatContext *s, AVPacket *pkt)
214{
215    WebMChunkContext *wc = s->priv_data;
216    AVFormatContext *oc = wc->avf;
217    AVStream *st = s->streams[pkt->stream_index];
218    int ret;
219
220    if (st->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
221        if (wc->prev_pts != AV_NOPTS_VALUE)
222            wc->duration_written += av_rescale_q(pkt->pts - wc->prev_pts,
223                                                 st->time_base,
224                                                 (AVRational) {1, 1000});
225        wc->prev_pts = pkt->pts;
226    }
227
228    // For video, a new chunk is started only on key frames. For audio, a new
229    // chunk is started based on chunk_duration. Also, a new chunk is started
230    // unconditionally if there is no currently open chunk.
231    if (!oc->pb || (st->codecpar->codec_type == AVMEDIA_TYPE_VIDEO &&
232         (pkt->flags & AV_PKT_FLAG_KEY)) ||
233        (st->codecpar->codec_type == AVMEDIA_TYPE_AUDIO &&
234         wc->duration_written >= wc->chunk_duration)) {
235        wc->duration_written = 0;
236        if ((ret = chunk_end(s, 1)) < 0 || (ret = chunk_start(s)) < 0) {
237            return ret;
238        }
239    }
240
241    // We only have one stream, so use the non-interleaving av_write_frame.
242    return av_write_frame(oc, pkt);
243}
244
245static int webm_chunk_write_trailer(AVFormatContext *s)
246{
247    WebMChunkContext *wc = s->priv_data;
248    AVFormatContext *oc = wc->avf;
249    int ret;
250
251    if (!oc->pb) {
252        ret = chunk_start(s);
253        if (ret < 0)
254            return ret;
255    }
256    ret = av_write_trailer(oc);
257    if (ret < 0)
258        return ret;
259    return chunk_end(s, 0);
260}
261
262static void webm_chunk_deinit(AVFormatContext *s)
263{
264    WebMChunkContext *wc = s->priv_data;
265
266    if (!wc->avf)
267        return;
268
269    if (wc->header_written)
270        ffio_free_dyn_buf(&wc->avf->pb);
271    else
272        ff_format_io_close(s, &wc->avf->pb);
273    avformat_free_context(wc->avf);
274    wc->avf = NULL;
275}
276
277#define OFFSET(x) offsetof(WebMChunkContext, x)
278static const AVOption options[] = {
279    { "chunk_start_index",  "start index of the chunk", OFFSET(chunk_index), AV_OPT_TYPE_INT, {.i64 = 0}, 0, INT_MAX, AV_OPT_FLAG_ENCODING_PARAM },
280    { "header", "filename of the header where the initialization data will be written", OFFSET(header_filename), AV_OPT_TYPE_STRING, {.str = NULL}, 0, 0, AV_OPT_FLAG_ENCODING_PARAM },
281    { "audio_chunk_duration", "duration of each chunk in milliseconds", OFFSET(chunk_duration), AV_OPT_TYPE_INT, {.i64 = 5000}, 0, INT_MAX, AV_OPT_FLAG_ENCODING_PARAM },
282    { "method", "set the HTTP method", OFFSET(http_method), AV_OPT_TYPE_STRING, {.str = NULL},  0, 0, AV_OPT_FLAG_ENCODING_PARAM },
283    { NULL },
284};
285
286static const AVClass webm_chunk_class = {
287    .class_name = "WebM Chunk Muxer",
288    .item_name  = av_default_item_name,
289    .option     = options,
290    .version    = LIBAVUTIL_VERSION_INT,
291};
292
293const AVOutputFormat ff_webm_chunk_muxer = {
294    .name           = "webm_chunk",
295    .long_name      = NULL_IF_CONFIG_SMALL("WebM Chunk Muxer"),
296    .mime_type      = "video/webm",
297    .extensions     = "chk",
298    .flags          = AVFMT_NOFILE | AVFMT_GLOBALHEADER | AVFMT_NEEDNUMBER |
299                      AVFMT_TS_NONSTRICT,
300    .priv_data_size = sizeof(WebMChunkContext),
301    .init           = webm_chunk_init,
302    .write_header   = webm_chunk_write_header,
303    .write_packet   = webm_chunk_write_packet,
304    .write_trailer  = webm_chunk_write_trailer,
305    .deinit         = webm_chunk_deinit,
306    .priv_class     = &webm_chunk_class,
307};
308