Avformat mux: fix a memory leak
[vlc.git] / modules / demux / avformat / mux.c
1 /*****************************************************************************
2  * mux.c: muxer using libavformat
3  *****************************************************************************
4  * Copyright (C) 2006 VLC authors and VideoLAN
5  * $Id$
6  *
7  * Authors: Gildas Bazin <gbazin@videolan.org>
8  *
9  * This program is free software; you can redistribute it and/or modify it
10  * under the terms of the GNU Lesser General Public License as published by
11  * the Free Software Foundation; either version 2.1 of the License, or
12  * (at your option) any later version.
13  *
14  * This program is distributed in the hope that it will be useful,
15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17  * GNU Lesser General Public License for more details.
18  *
19  * You should have received a copy of the GNU Lesser General Public License
20  * along with this program; if not, write to the Free Software Foundation,
21  * Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.
22  *****************************************************************************/
23
24 /*****************************************************************************
25  * Preamble
26  *****************************************************************************/
27
28 #ifdef HAVE_CONFIG_H
29 # include "config.h"
30 #endif
31
32 #include <vlc_common.h>
33 #include <vlc_block.h>
34 #include <vlc_sout.h>
35
36 #include <libavformat/avformat.h>
37
38 #include "avformat.h"
39 #include "../../codec/avcodec/avcodec.h"
40 #include "../../codec/avcodec/avcommon.h"
41 #include "../xiph.h"
42
43
44 //#define AVFORMAT_DEBUG 1
45
46 static const char *const ppsz_mux_options[] = {
47     "mux", "options", NULL
48 };
49
50 /*****************************************************************************
51  * mux_sys_t: mux descriptor
52  *****************************************************************************/
53 struct sout_mux_sys_t
54 {
55     AVIOContext     *io;
56     int             io_buffer_size;
57     uint8_t        *io_buffer;
58
59     AVFormatContext *oc;
60
61     bool     b_write_header;
62     bool     b_write_keyframe;
63     bool     b_error;
64 };
65
66 /*****************************************************************************
67  * Local prototypes
68  *****************************************************************************/
69 static int Control  ( sout_mux_t *, int, va_list );
70 static int AddStream( sout_mux_t *, sout_input_t * );
71 static int DelStream( sout_mux_t *, sout_input_t * );
72 static int Mux      ( sout_mux_t * );
73
74 static int IOWrite( void *opaque, uint8_t *buf, int buf_size );
75 static int64_t IOSeek( void *opaque, int64_t offset, int whence );
76
77 /*****************************************************************************
78  * Open
79  *****************************************************************************/
80 int OpenMux( vlc_object_t *p_this )
81 {
82     AVOutputFormat *file_oformat;
83     sout_mux_t *p_mux = (sout_mux_t*)p_this;
84     bool dummy = !strcmp( p_mux->p_access->psz_access, "dummy");
85
86     if( dummy && strlen(p_mux->p_access->psz_path)
87                               >= sizeof (((AVFormatContext *)NULL)->filename) )
88         return VLC_EGENERIC;
89
90     vlc_init_avformat(p_this);
91
92     config_ChainParse( p_mux, "sout-avformat-", ppsz_mux_options, p_mux->p_cfg );
93
94     /* Find the requested muxer */
95     char *psz_mux = var_InheritString( p_mux, "sout-avformat-mux" );
96     if( psz_mux )
97     {
98         file_oformat = av_guess_format( psz_mux, NULL, NULL );
99         free( psz_mux );
100     }
101     else
102     {
103         file_oformat =
104             av_guess_format( NULL, p_mux->p_access->psz_path, NULL);
105     }
106     if (!file_oformat)
107     {
108       msg_Err( p_mux, "unable for find a suitable output format" );
109       return VLC_EGENERIC;
110     }
111
112     sout_mux_sys_t *p_sys = malloc( sizeof( sout_mux_sys_t ) );
113     if( unlikely(p_sys == NULL) )
114         return VLC_ENOMEM;
115
116     p_mux->p_sys = p_sys;
117     p_sys->oc = avformat_alloc_context();
118     p_sys->oc->oformat = file_oformat;
119     /* If we use dummy access, let avformat write output */
120     if( dummy )
121         strcpy( p_sys->oc->filename, p_mux->p_access->psz_path );
122
123     /* Create I/O wrapper */
124     p_sys->io_buffer_size = 10 * 1024 * 1024;  /* FIXME */
125     p_sys->io_buffer = malloc( p_sys->io_buffer_size );
126
127     bool b_can_seek;
128     if( sout_AccessOutControl( p_mux->p_access, ACCESS_OUT_CAN_SEEK, &b_can_seek ) )
129         b_can_seek = false;
130     p_sys->io = avio_alloc_context(
131         p_sys->io_buffer, p_sys->io_buffer_size,
132         1, p_mux, NULL, IOWrite, b_can_seek ? IOSeek : NULL );
133
134     p_sys->oc->pb = p_sys->io;
135     p_sys->oc->nb_streams = 0;
136
137     p_sys->b_write_header = true;
138     p_sys->b_write_keyframe = false;
139     p_sys->b_error = false;
140
141     /* Fill p_mux fields */
142     p_mux->pf_control   = Control;
143     p_mux->pf_addstream = AddStream;
144     p_mux->pf_delstream = DelStream;
145     p_mux->pf_mux       = Mux;
146
147     return VLC_SUCCESS;
148 }
149
150 /*****************************************************************************
151  * Close
152  *****************************************************************************/
153 void CloseMux( vlc_object_t *p_this )
154 {
155     sout_mux_t *p_mux = (sout_mux_t*)p_this;
156     sout_mux_sys_t *p_sys = p_mux->p_sys;
157
158     if( !p_sys->b_write_header && !p_sys->b_error && av_write_trailer( p_sys->oc ) < 0 )
159     {
160         msg_Err( p_mux, "could not write trailer" );
161     }
162
163     avformat_free_context(p_sys->oc);
164     av_free(p_sys->io);
165
166     free( p_sys->io_buffer );
167     free( p_sys );
168 }
169
170 /*****************************************************************************
171  * AddStream
172  *****************************************************************************/
173 static int AddStream( sout_mux_t *p_mux, sout_input_t *p_input )
174 {
175     sout_mux_sys_t *p_sys = p_mux->p_sys;
176     es_format_t *fmt = p_input->p_fmt;
177     AVCodecContext *codec;
178     AVStream *stream;
179     unsigned i_codec_id;
180
181     msg_Dbg( p_mux, "adding input" );
182
183     if( !GetFfmpegCodec( fmt->i_codec, 0, &i_codec_id, 0 ) )
184     {
185         msg_Dbg( p_mux, "couldn't find codec for fourcc '%4.4s'",
186                  (char *)&fmt->i_codec );
187         return VLC_EGENERIC;
188     }
189
190     unsigned opus_size[XIPH_MAX_HEADER_COUNT];
191     void     *opus_packet[XIPH_MAX_HEADER_COUNT];
192     if( fmt->i_codec == VLC_CODEC_OPUS )
193     {
194         unsigned count;
195         /* Only transmits the first packet (OpusHead) */
196         if( xiph_SplitHeaders(opus_size, opus_packet, &count, fmt->i_extra, fmt->p_extra ) ) {
197             count = 0;
198         }
199         if (count != 2 || opus_size[0] < 19) {
200             msg_Err(p_mux, "Invalid Opus header");
201             return VLC_EGENERIC;
202         }
203     }
204
205     p_input->p_sys = malloc( sizeof( int ) );
206     *((int *)p_input->p_sys) = p_sys->oc->nb_streams;
207
208     if( fmt->i_cat != VIDEO_ES && fmt->i_cat != AUDIO_ES)
209     {
210         msg_Warn( p_mux, "Unhandled ES category" );
211         return VLC_EGENERIC;
212     }
213
214     stream = avformat_new_stream( p_sys->oc, NULL);
215     if( !stream )
216     {
217         free( p_input->p_sys );
218         return VLC_EGENERIC;
219     }
220     codec = stream->codec;
221
222     codec->opaque = p_mux;
223
224     switch( fmt->i_cat )
225     {
226     case AUDIO_ES:
227         codec->codec_type = AVMEDIA_TYPE_AUDIO;
228         codec->channels = fmt->audio.i_channels;
229         codec->sample_rate = fmt->audio.i_rate;
230         codec->time_base = (AVRational){1, codec->sample_rate};
231         codec->frame_size = fmt->audio.i_frame_length;
232         if (fmt->i_bitrate == 0) {
233             msg_Warn( p_mux, "Missing audio bitrate, assuming 64k" );
234             fmt->i_bitrate = 64000;
235         }
236         break;
237
238     case VIDEO_ES:
239         if( !fmt->video.i_frame_rate || !fmt->video.i_frame_rate_base ) {
240             msg_Warn( p_mux, "Missing frame rate, assuming 25fps" );
241             fmt->video.i_frame_rate = 25;
242             fmt->video.i_frame_rate_base = 1;
243         } else
244             msg_Dbg( p_mux, "Muxing framerate will be %d/%d = %.2f fps",
245                     fmt->video.i_frame_rate,
246                     fmt->video.i_frame_rate_base,
247                     (double)fmt->video.i_frame_rate/(double)fmt->video.i_frame_rate_base );
248
249         codec->codec_type = AVMEDIA_TYPE_VIDEO;
250         codec->width = fmt->video.i_width;
251         codec->height = fmt->video.i_height;
252         av_reduce( &codec->sample_aspect_ratio.num,
253                    &codec->sample_aspect_ratio.den,
254                    fmt->video.i_sar_num,
255                    fmt->video.i_sar_den, 1 << 30 /* something big */ );
256         msg_Dbg(p_mux, "Muxing aspect ratio will be %d/%d",
257                 fmt->video.i_sar_num, fmt->video.i_sar_den);
258         stream->sample_aspect_ratio.den = codec->sample_aspect_ratio.den;
259         stream->sample_aspect_ratio.num = codec->sample_aspect_ratio.num;
260         codec->time_base.den = fmt->video.i_frame_rate;
261         codec->time_base.num = fmt->video.i_frame_rate_base;
262         if (fmt->i_bitrate == 0) {
263             msg_Warn( p_mux, "Missing video bitrate, assuming 512k" );
264             fmt->i_bitrate = 512000;
265         } else
266             msg_Dbg( p_mux, "Muxing video bitrate will be %d", fmt->i_bitrate );
267         break;
268
269     }
270
271     codec->bit_rate = fmt->i_bitrate;
272     codec->codec_tag = av_codec_get_tag( p_sys->oc->oformat->codec_tag, i_codec_id );
273     if( !codec->codec_tag && i_codec_id == AV_CODEC_ID_MP2 )
274     {
275         i_codec_id = AV_CODEC_ID_MP3;
276         codec->codec_tag = av_codec_get_tag( p_sys->oc->oformat->codec_tag, i_codec_id );
277     }
278     codec->codec_id = i_codec_id;
279
280     if( fmt->i_extra )
281     {
282         if( fmt->i_codec == VLC_CODEC_OPUS )
283         {
284             codec->extradata_size = opus_size[0];
285             codec->extradata = av_malloc( opus_size[0] );
286             memcpy( codec->extradata, opus_packet[0], opus_size[0] );
287         }
288         else
289         {
290             codec->extradata_size = fmt->i_extra;
291             codec->extradata = av_malloc( fmt->i_extra );
292             memcpy( codec->extradata, fmt->p_extra, fmt->i_extra );
293         }
294     }
295
296     return VLC_SUCCESS;
297 }
298
299 /*****************************************************************************
300  * DelStream
301  *****************************************************************************/
302 static int DelStream( sout_mux_t *p_mux, sout_input_t *p_input )
303 {
304     msg_Dbg( p_mux, "removing input" );
305     free( p_input->p_sys );
306     return VLC_SUCCESS;
307 }
308
309 static int MuxBlock( sout_mux_t *p_mux, sout_input_t *p_input )
310 {
311     sout_mux_sys_t *p_sys = p_mux->p_sys;
312     block_t *p_data = block_FifoGet( p_input->p_fifo );
313     int i_stream = *((int *)p_input->p_sys);
314     AVStream *p_stream = p_sys->oc->streams[i_stream];
315     AVPacket pkt;
316
317     memset( &pkt, 0, sizeof(AVPacket) );
318
319     av_init_packet(&pkt);
320     pkt.data = p_data->p_buffer;
321     pkt.size = p_data->i_buffer;
322     pkt.stream_index = i_stream;
323
324     if( p_data->i_flags & BLOCK_FLAG_TYPE_I )
325     {
326 #ifdef AVFMT_ALLOW_FLUSH
327         /* Make sure we don't inadvertedly mark buffered data as keyframes. */
328         if( p_sys->oc->oformat->flags & AVFMT_ALLOW_FLUSH )
329             av_write_frame( p_sys->oc, NULL );
330 #endif
331
332         p_sys->b_write_keyframe = true;
333         pkt.flags |= AV_PKT_FLAG_KEY;
334     }
335
336     if( p_data->i_pts > 0 )
337         pkt.pts = p_data->i_pts * p_stream->time_base.den /
338             CLOCK_FREQ / p_stream->time_base.num;
339     if( p_data->i_dts > 0 )
340         pkt.dts = p_data->i_dts * p_stream->time_base.den /
341             CLOCK_FREQ / p_stream->time_base.num;
342
343     /* this is another hack to prevent libavformat from triggering the "non monotone timestamps" check in avformat/utils.c */
344     p_stream->cur_dts = ( p_data->i_dts * p_stream->time_base.den /
345             CLOCK_FREQ / p_stream->time_base.num ) - 1;
346
347     if( av_write_frame( p_sys->oc, &pkt ) < 0 )
348     {
349         msg_Err( p_mux, "could not write frame (pts: %"PRId64", dts: %"PRId64") "
350                  "(pkt pts: %"PRId64", dts: %"PRId64")",
351                  p_data->i_pts, p_data->i_dts, pkt.pts, pkt.dts );
352         block_Release( p_data );
353         return VLC_EGENERIC;
354     }
355
356     block_Release( p_data );
357     return VLC_SUCCESS;
358 }
359
360 /*****************************************************************************
361  * Mux: multiplex available data in input fifos
362  *****************************************************************************/
363 static int Mux( sout_mux_t *p_mux )
364 {
365     sout_mux_sys_t *p_sys = p_mux->p_sys;
366
367     if( p_sys->b_error ) return VLC_EGENERIC;
368
369     if( p_sys->b_write_header )
370     {
371         int error;
372         msg_Dbg( p_mux, "writing header" );
373
374         char *psz_opts = var_GetNonEmptyString( p_mux, "sout-avformat-options" );
375         AVDictionary *options = NULL;
376         if (psz_opts && *psz_opts)
377             options = vlc_av_get_options(psz_opts);
378         free(psz_opts);
379         error = avformat_write_header( p_sys->oc, options ? &options : NULL);
380         AVDictionaryEntry *t = NULL;
381         while ((t = av_dict_get(options, "", t, AV_DICT_IGNORE_SUFFIX))) {
382             msg_Err( p_mux, "Unknown option \"%s\"", t->key );
383         }
384         av_dict_free(&options);
385         if( error < 0 )
386         {
387             msg_Err( p_mux, "could not write header: %s",
388                      vlc_strerror_c(AVUNERROR(error)) );
389             p_sys->b_write_header = false;
390             p_sys->b_error = true;
391             return VLC_EGENERIC;
392         }
393
394         avio_flush( p_sys->oc->pb );
395         p_sys->b_write_header = false;
396     }
397
398     for( ;; )
399     {
400         mtime_t i_dts;
401
402         int i_stream = sout_MuxGetStream( p_mux, 1, &i_dts );
403         if( i_stream < 0 )
404             return VLC_SUCCESS;
405
406         MuxBlock( p_mux, p_mux->pp_inputs[i_stream] );
407     }
408
409     return VLC_SUCCESS;
410 }
411
412 /*****************************************************************************
413  * Control:
414  *****************************************************************************/
415 static int Control( sout_mux_t *p_mux, int i_query, va_list args )
416 {
417     bool *pb_bool;
418
419     switch( i_query )
420     {
421     case MUX_CAN_ADD_STREAM_WHILE_MUXING:
422         pb_bool = (bool*)va_arg( args, bool * );
423         *pb_bool = false;
424         return VLC_SUCCESS;
425
426     case MUX_GET_ADD_STREAM_WAIT:
427         pb_bool = (bool*)va_arg( args, bool * );
428         *pb_bool = true;
429         return VLC_SUCCESS;
430
431     case MUX_GET_MIME:
432     {
433         char **ppsz = (char**)va_arg( args, char ** );
434         *ppsz = strdup( p_mux->p_sys->oc->oformat->mime_type );
435         return VLC_SUCCESS;
436     }
437
438     default:
439         return VLC_EGENERIC;
440     }
441 }
442
443 /*****************************************************************************
444  * I/O wrappers for libavformat
445  *****************************************************************************/
446 static int IOWrite( void *opaque, uint8_t *buf, int buf_size )
447 {
448     sout_mux_t *p_mux = opaque;
449     sout_mux_sys_t *p_sys = p_mux->p_sys;
450     int i_ret;
451
452 #ifdef AVFORMAT_DEBUG
453     msg_Dbg( p_mux, "IOWrite %i bytes", buf_size );
454 #endif
455
456     block_t *p_buf = block_Alloc( buf_size );
457     if( buf_size > 0 ) memcpy( p_buf->p_buffer, buf, buf_size );
458
459     if( p_sys->b_write_header )
460         p_buf->i_flags |= BLOCK_FLAG_HEADER;
461
462     if( p_sys->b_write_keyframe )
463     {
464         p_buf->i_flags |= BLOCK_FLAG_TYPE_I;
465         p_sys->b_write_keyframe = false;
466     }
467
468     i_ret = sout_AccessOutWrite( p_mux->p_access, p_buf );
469     return i_ret ? i_ret : -1;
470 }
471
472 static int64_t IOSeek( void *opaque, int64_t offset, int whence )
473 {
474     sout_mux_t *p_mux = opaque;
475
476 #ifdef AVFORMAT_DEBUG
477     msg_Dbg( p_mux, "IOSeek offset: %"PRId64", whence: %i", offset, whence );
478 #endif
479
480     switch( whence )
481     {
482     case SEEK_SET:
483         return sout_AccessOutSeek( p_mux->p_access, offset );
484     case SEEK_CUR:
485     case SEEK_END:
486     default:
487         return -1;
488     }
489 }