matroskaenc.c 31.2 KB
Newer Older
1
/*
David Conrad's avatar
David Conrad committed
2
 * Matroska muxer
3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
 * Copyright (c) 2007 David Conrad
 *
 * This file is part of FFmpeg.
 *
 * FFmpeg is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * FFmpeg is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with FFmpeg; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 */

#include "avformat.h"
#include "riff.h"
24
#include "isom.h"
25
#include "matroska.h"
26
#include "avc.h"
27
#include "flacenc.h"
28
#include "libavutil/intreadwrite.h"
29 30
#include "libavutil/md5.h"
#include "libavcodec/xiph.h"
31
#include "libavcodec/mpeg4audio.h"
32

33
typedef struct ebml_master {
34
    int64_t         pos;                ///< absolute offset in the file where the master's elements start
35 36 37
    int             sizebytes;          ///< how many bytes were reserved for the size
} ebml_master;

38 39 40 41 42 43
typedef struct mkv_seekhead_entry {
    unsigned int    elementid;
    uint64_t        segmentpos;
} mkv_seekhead_entry;

typedef struct mkv_seekhead {
44 45
    int64_t                 filepos;
    int64_t                 segment_offset;     ///< the file offset to the beginning of the segment
David Conrad's avatar
David Conrad committed
46
    int                     reserved_size;      ///< -1 if appending to file
47 48 49 50 51
    int                     max_entries;
    mkv_seekhead_entry      *entries;
    int                     num_entries;
} mkv_seekhead;

David Conrad's avatar
David Conrad committed
52 53 54
typedef struct {
    uint64_t        pts;
    int             tracknum;
55
    int64_t         cluster_pos;        ///< file offset of the cluster containing the block
David Conrad's avatar
David Conrad committed
56 57 58
} mkv_cuepoint;

typedef struct {
59
    int64_t         segment_offset;
David Conrad's avatar
David Conrad committed
60 61 62 63
    mkv_cuepoint    *entries;
    int             num_entries;
} mkv_cues;

64
typedef struct MatroskaMuxContext {
65
    ebml_master     segment;
66 67
    int64_t         segment_offset;
    int64_t         segment_uid;
68
    ebml_master     cluster;
69
    int64_t         cluster_pos;        ///< file offset of the current cluster
David Conrad's avatar
David Conrad committed
70
    uint64_t        cluster_pts;
71
    int64_t         duration_offset;
David Conrad's avatar
David Conrad committed
72
    uint64_t        duration;
73 74
    mkv_seekhead    *main_seekhead;
    mkv_seekhead    *cluster_seekhead;
David Conrad's avatar
David Conrad committed
75
    mkv_cues        *cues;
76 77

    struct AVMD5    *md5_ctx;
78 79
} MatroskaMuxContext;

80

David Conrad's avatar
David Conrad committed
81 82
/** 2 bytes * 3 for EBML IDs, 3 1-byte EBML lengths, 8 bytes for 64 bit
 * offset, 4 bytes for target EBML ID */
83 84
#define MAX_SEEKENTRY_SIZE 21

David Conrad's avatar
David Conrad committed
85 86
/** per-cuepoint-track - 3 1-byte EBML IDs, 3 1-byte EBML sizes, 2
 * 8-byte uint max */
87 88
#define MAX_CUETRACKPOS_SIZE 22

David Conrad's avatar
David Conrad committed
89
/** per-cuepoint - 2 1-byte EBML IDs, 2 1-byte EBML sizes, 8-byte uint max */
90 91 92
#define MAX_CUEPOINT_SIZE(num_tracks) 12 + MAX_CUETRACKPOS_SIZE*num_tracks


David Conrad's avatar
David Conrad committed
93 94 95 96 97
static int ebml_id_size(unsigned int id)
{
    return (av_log2(id+1)-1)/7+1;
}

98 99
static void put_ebml_id(ByteIOContext *pb, unsigned int id)
{
David Conrad's avatar
David Conrad committed
100 101 102
    int i = ebml_id_size(id);
    while (i--)
        put_byte(pb, id >> (i*8));
103 104
}

105
/**
106
 * Write an EBML size meaning "unknown size".
107
 *
108
 * @param bytes The number of bytes the size should occupy (maximum: 8).
109 110 111
 */
static void put_ebml_size_unknown(ByteIOContext *pb, int bytes)
{
112
    assert(bytes <= 8);
David Conrad's avatar
David Conrad committed
113 114 115
    put_byte(pb, 0x1ff >> bytes);
    while (--bytes)
        put_byte(pb, 0xff);
116 117
}

118
/**
119
 * Calculate how many bytes are needed to represent a given number in EBML.
120
 */
121
static int ebml_num_size(uint64_t num)
122 123
{
    int bytes = 1;
124
    while ((num+1) >> bytes*7) bytes++;
125 126 127
    return bytes;
}

128
/**
129
 * Write a number in EBML variable length format.
130
 *
131
 * @param bytes The number of bytes that need to be used to write the number.
132 133
 *              If zero, any number of bytes can be used.
 */
134
static void put_ebml_num(ByteIOContext *pb, uint64_t num, int bytes)
135
{
136
    int i, needed_bytes = ebml_num_size(num);
137 138

    // sizes larger than this are currently undefined in EBML
139
    assert(num < (1ULL<<56)-1);
140

141 142 143
    if (bytes == 0)
        // don't care how many bytes are used, so use the min
        bytes = needed_bytes;
David Conrad's avatar
David Conrad committed
144 145
    // the bytes needed to write the given size would exceed the bytes
    // that we need to use, so write unknown size. This shouldn't happen.
146
    assert(bytes >= needed_bytes);
147

148
    num |= 1ULL << bytes*7;
David Conrad's avatar
David Conrad committed
149
    for (i = bytes - 1; i >= 0; i--)
150
        put_byte(pb, num >> i*8);
151 152 153 154
}

static void put_ebml_uint(ByteIOContext *pb, unsigned int elementid, uint64_t val)
{
155
    int i, bytes = 1;
156 157
    uint64_t tmp = val;
    while (tmp>>=8) bytes++;
158 159

    put_ebml_id(pb, elementid);
160
    put_ebml_num(pb, bytes, 0);
161 162
    for (i = bytes - 1; i >= 0; i--)
        put_byte(pb, val >> i*8);
163 164
}

David Conrad's avatar
David Conrad committed
165 166 167
static void put_ebml_float(ByteIOContext *pb, unsigned int elementid, double val)
{
    put_ebml_id(pb, elementid);
168
    put_ebml_num(pb, 8, 0);
David Conrad's avatar
David Conrad committed
169 170 171
    put_be64(pb, av_dbl2int(val));
}

172
static void put_ebml_binary(ByteIOContext *pb, unsigned int elementid,
David Conrad's avatar
David Conrad committed
173
                            const uint8_t *buf, int size)
174 175
{
    put_ebml_id(pb, elementid);
176
    put_ebml_num(pb, size, 0);
177 178 179
    put_buffer(pb, buf, size);
}

David Conrad's avatar
David Conrad committed
180
static void put_ebml_string(ByteIOContext *pb, unsigned int elementid, const char *str)
181 182 183 184
{
    put_ebml_binary(pb, elementid, str, strlen(str));
}

David Conrad's avatar
David Conrad committed
185
/**
David Conrad's avatar
David Conrad committed
186 187
 * Writes a void element of a given size. Useful for reserving space in
 * the file to be written to later.
David Conrad's avatar
David Conrad committed
188
 *
David Conrad's avatar
David Conrad committed
189
 * @param size The number of bytes to reserve, which must be at least 2.
David Conrad's avatar
David Conrad committed
190
 */
David Conrad's avatar
David Conrad committed
191 192
static void put_ebml_void(ByteIOContext *pb, uint64_t size)
{
193
    int64_t currentpos = url_ftell(pb);
David Conrad's avatar
David Conrad committed
194

195
    assert(size >= 2);
David Conrad's avatar
David Conrad committed
196 197

    put_ebml_id(pb, EBML_ID_VOID);
David Conrad's avatar
David Conrad committed
198 199 200
    // we need to subtract the length needed to store the size from the
    // size we need to reserve so 2 cases, we use 8 bytes to store the
    // size if possible, 1 byte otherwise
David Conrad's avatar
David Conrad committed
201
    if (size < 10)
202
        put_ebml_num(pb, size-1, 0);
David Conrad's avatar
David Conrad committed
203
    else
204
        put_ebml_num(pb, size-9, 8);
205 206
    while(url_ftell(pb) < currentpos + size)
        put_byte(pb, 0);
David Conrad's avatar
David Conrad committed
207 208
}

209
static ebml_master start_ebml_master(ByteIOContext *pb, unsigned int elementid, uint64_t expectedsize)
210
{
211
    int bytes = expectedsize ? ebml_num_size(expectedsize) : 8;
212
    put_ebml_id(pb, elementid);
213 214
    put_ebml_size_unknown(pb, bytes);
    return (ebml_master){ url_ftell(pb), bytes };
215 216
}

217
static void end_ebml_master(ByteIOContext *pb, ebml_master master)
218
{
219
    int64_t pos = url_ftell(pb);
220

David Conrad's avatar
David Conrad committed
221 222 223 224
    // leave the unknown size for masters when streaming
    if (url_is_streamed(pb))
        return;

225
    url_fseek(pb, master.pos - master.sizebytes, SEEK_SET);
226
    put_ebml_num(pb, pos - master.pos, master.sizebytes);
227 228 229
    url_fseek(pb, pos, SEEK_SET);
}

230 231 232 233 234 235 236 237
static void put_xiph_size(ByteIOContext *pb, int size)
{
    int i;
    for (i = 0; i < size / 255; i++)
        put_byte(pb, 255);
    put_byte(pb, size % 255);
}

David Conrad's avatar
David Conrad committed
238
/**
David Conrad's avatar
David Conrad committed
239 240 241 242
 * Initialize a mkv_seekhead element to be ready to index level 1 Matroska
 * elements. If a maximum number of elements is specified, enough space
 * will be reserved at the current file location to write a seek head of
 * that size.
David Conrad's avatar
David Conrad committed
243
 *
David Conrad's avatar
David Conrad committed
244
 * @param segment_offset The absolute offset to the position in the file
David Conrad's avatar
David Conrad committed
245 246
 *                       where the segment begins.
 * @param numelements The maximum number of elements that will be indexed
David Conrad's avatar
David Conrad committed
247
 *                    by this seek head, 0 if unlimited.
David Conrad's avatar
David Conrad committed
248
 */
249
static mkv_seekhead * mkv_start_seekhead(ByteIOContext *pb, int64_t segment_offset, int numelements)
250 251 252 253 254 255 256 257 258
{
    mkv_seekhead *new_seekhead = av_mallocz(sizeof(mkv_seekhead));
    if (new_seekhead == NULL)
        return NULL;

    new_seekhead->segment_offset = segment_offset;

    if (numelements > 0) {
        new_seekhead->filepos = url_ftell(pb);
David Conrad's avatar
David Conrad committed
259 260 261
        // 21 bytes max for a seek entry, 10 bytes max for the SeekHead ID
        // and size, and 3 bytes to guarantee that an EBML void element
        // will fit afterwards
262
        new_seekhead->reserved_size = numelements * MAX_SEEKENTRY_SIZE + 13;
263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278
        new_seekhead->max_entries = numelements;
        put_ebml_void(pb, new_seekhead->reserved_size);
    }
    return new_seekhead;
}

static int mkv_add_seekhead_entry(mkv_seekhead *seekhead, unsigned int elementid, uint64_t filepos)
{
    mkv_seekhead_entry *entries = seekhead->entries;

    // don't store more elements than we reserved space for
    if (seekhead->max_entries > 0 && seekhead->max_entries <= seekhead->num_entries)
        return -1;

    entries = av_realloc(entries, (seekhead->num_entries + 1) * sizeof(mkv_seekhead_entry));
    if (entries == NULL)
279
        return AVERROR(ENOMEM);
280

281 282
    entries[seekhead->num_entries  ].elementid = elementid;
    entries[seekhead->num_entries++].segmentpos = filepos - seekhead->segment_offset;
283 284 285 286 287

    seekhead->entries = entries;
    return 0;
}

David Conrad's avatar
David Conrad committed
288
/**
David Conrad's avatar
David Conrad committed
289 290 291 292
 * Write the seek head to the file and free it. If a maximum number of
 * elements was specified to mkv_start_seekhead(), the seek head will
 * be written at the location reserved for it. Otherwise, it is written
 * at the current location in the file.
David Conrad's avatar
David Conrad committed
293
 *
David Conrad's avatar
David Conrad committed
294
 * @return The file offset where the seekhead was written.
David Conrad's avatar
David Conrad committed
295
 */
296
static int64_t mkv_write_seekhead(ByteIOContext *pb, mkv_seekhead *seekhead)
297
{
298
    ebml_master metaseek, seekentry;
299
    int64_t currentpos;
300 301 302 303 304 305 306
    int i;

    currentpos = url_ftell(pb);

    if (seekhead->reserved_size > 0)
        url_fseek(pb, seekhead->filepos, SEEK_SET);

307
    metaseek = start_ebml_master(pb, MATROSKA_ID_SEEKHEAD, seekhead->reserved_size);
308 309 310
    for (i = 0; i < seekhead->num_entries; i++) {
        mkv_seekhead_entry *entry = &seekhead->entries[i];

311
        seekentry = start_ebml_master(pb, MATROSKA_ID_SEEKENTRY, MAX_SEEKENTRY_SIZE);
312 313

        put_ebml_id(pb, MATROSKA_ID_SEEKID);
314
        put_ebml_num(pb, ebml_id_size(entry->elementid), 0);
315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334
        put_ebml_id(pb, entry->elementid);

        put_ebml_uint(pb, MATROSKA_ID_SEEKPOSITION, entry->segmentpos);
        end_ebml_master(pb, seekentry);
    }
    end_ebml_master(pb, metaseek);

    if (seekhead->reserved_size > 0) {
        uint64_t remaining = seekhead->filepos + seekhead->reserved_size - url_ftell(pb);
        put_ebml_void(pb, remaining);
        url_fseek(pb, currentpos, SEEK_SET);

        currentpos = seekhead->filepos;
    }
    av_free(seekhead->entries);
    av_free(seekhead);

    return currentpos;
}

335
static mkv_cues * mkv_start_cues(int64_t segment_offset)
David Conrad's avatar
David Conrad committed
336 337 338 339 340 341 342 343 344
{
    mkv_cues *cues = av_mallocz(sizeof(mkv_cues));
    if (cues == NULL)
        return NULL;

    cues->segment_offset = segment_offset;
    return cues;
}

345
static int mkv_add_cuepoint(mkv_cues *cues, AVPacket *pkt, int64_t cluster_pos)
David Conrad's avatar
David Conrad committed
346 347 348 349 350
{
    mkv_cuepoint *entries = cues->entries;

    entries = av_realloc(entries, (cues->num_entries + 1) * sizeof(mkv_cuepoint));
    if (entries == NULL)
351
        return AVERROR(ENOMEM);
David Conrad's avatar
David Conrad committed
352

353 354 355
    entries[cues->num_entries  ].pts = pkt->pts;
    entries[cues->num_entries  ].tracknum = pkt->stream_index + 1;
    entries[cues->num_entries++].cluster_pos = cluster_pos - cues->segment_offset;
David Conrad's avatar
David Conrad committed
356 357 358 359 360

    cues->entries = entries;
    return 0;
}

361
static int64_t mkv_write_cues(ByteIOContext *pb, mkv_cues *cues, int num_tracks)
David Conrad's avatar
David Conrad committed
362
{
363
    ebml_master cues_element;
364
    int64_t currentpos;
David Conrad's avatar
David Conrad committed
365 366 367
    int i, j;

    currentpos = url_ftell(pb);
368
    cues_element = start_ebml_master(pb, MATROSKA_ID_CUES, 0);
David Conrad's avatar
David Conrad committed
369 370

    for (i = 0; i < cues->num_entries; i++) {
371
        ebml_master cuepoint, track_positions;
David Conrad's avatar
David Conrad committed
372 373 374
        mkv_cuepoint *entry = &cues->entries[i];
        uint64_t pts = entry->pts;

375
        cuepoint = start_ebml_master(pb, MATROSKA_ID_POINTENTRY, MAX_CUEPOINT_SIZE(num_tracks));
David Conrad's avatar
David Conrad committed
376 377 378 379 380
        put_ebml_uint(pb, MATROSKA_ID_CUETIME, pts);

        // put all the entries from different tracks that have the exact same
        // timestamp into the same CuePoint
        for (j = 0; j < cues->num_entries - i && entry[j].pts == pts; j++) {
381
            track_positions = start_ebml_master(pb, MATROSKA_ID_CUETRACKPOSITION, MAX_CUETRACKPOS_SIZE);
David Conrad's avatar
David Conrad committed
382 383 384 385 386 387 388 389 390 391 392 393 394 395
            put_ebml_uint(pb, MATROSKA_ID_CUETRACK          , entry[j].tracknum   );
            put_ebml_uint(pb, MATROSKA_ID_CUECLUSTERPOSITION, entry[j].cluster_pos);
            end_ebml_master(pb, track_positions);
        }
        i += j - 1;
        end_ebml_master(pb, cuepoint);
    }
    end_ebml_master(pb, cues_element);

    av_free(cues->entries);
    av_free(cues);
    return currentpos;
}

396
static int put_xiph_codecpriv(AVFormatContext *s, ByteIOContext *pb, AVCodecContext *codec)
397 398 399 400
{
    uint8_t *header_start[3];
    int header_len[3];
    int first_header_size;
401
    int j;
402 403 404 405 406 407 408 409

    if (codec->codec_id == CODEC_ID_VORBIS)
        first_header_size = 30;
    else
        first_header_size = 42;

    if (ff_split_xiph_headers(codec->extradata, codec->extradata_size,
                              first_header_size, header_start, header_len) < 0) {
410
        av_log(s, AV_LOG_ERROR, "Extradata corrupt.\n");
411 412 413 414 415
        return -1;
    }

    put_byte(pb, 2);                    // number packets - 1
    for (j = 0; j < 2; j++) {
416
        put_xiph_size(pb, header_len[j]);
417 418 419
    }
    for (j = 0; j < 3; j++)
        put_buffer(pb, header_start[j], header_len[j]);
420 421

    return 0;
422
}
423

424
static void get_aac_sample_rates(AVFormatContext *s, AVCodecContext *codec, int *sample_rate, int *output_sample_rate)
425 426 427 428
{
    int sri;

    if (codec->extradata_size < 2) {
David Conrad's avatar
David Conrad committed
429
        av_log(s, AV_LOG_WARNING, "No AAC extradata, unable to determine samplerate.\n");
430 431 432 433 434
        return;
    }

    sri = ((codec->extradata[0] << 1) & 0xE) | (codec->extradata[1] >> 7);
    if (sri > 12) {
435
        av_log(s, AV_LOG_WARNING, "AAC samplerate index out of bounds\n");
436 437
        return;
    }
438
    *sample_rate = ff_mpeg4audio_sample_rates[sri];
439 440 441 442 443

    // if sbr, get output sample rate as well
    if (codec->extradata_size == 5) {
        sri = (codec->extradata[4] >> 3) & 0xF;
        if (sri > 12) {
444
            av_log(s, AV_LOG_WARNING, "AAC output samplerate index out of bounds\n");
445 446
            return;
        }
447
        *output_sample_rate = ff_mpeg4audio_sample_rates[sri];
448 449 450
    }
}

451
static int mkv_write_codecprivate(AVFormatContext *s, ByteIOContext *pb, AVCodecContext *codec, int native_id, int qt_id)
452
{
453
    ByteIOContext *dyn_cp;
454
    uint8_t *codecpriv;
455
    int ret, codecpriv_size;
456

457 458 459
    ret = url_open_dyn_buf(&dyn_cp);
    if(ret < 0)
        return ret;
460 461

    if (native_id) {
David Conrad's avatar
David Conrad committed
462
        if (codec->codec_id == CODEC_ID_VORBIS || codec->codec_id == CODEC_ID_THEORA)
463
            ret = put_xiph_codecpriv(s, dyn_cp, codec);
David Conrad's avatar
David Conrad committed
464
        else if (codec->codec_id == CODEC_ID_FLAC)
465
            ret = ff_flac_write_header(dyn_cp, codec);
466
        else if (codec->codec_id == CODEC_ID_H264)
467
            ret = ff_isom_write_avcc(dyn_cp, codec->extradata, codec->extradata_size);
David Conrad's avatar
David Conrad committed
468
        else if (codec->extradata_size)
469
            put_buffer(dyn_cp, codec->extradata, codec->extradata_size);
470
    } else if (codec->codec_type == CODEC_TYPE_VIDEO) {
471 472 473 474 475 476
        if (qt_id) {
            if (!codec->codec_tag)
                codec->codec_tag = codec_get_tag(codec_movvideo_tags, codec->codec_id);
            if (codec->extradata_size)
                put_buffer(dyn_cp, codec->extradata, codec->extradata_size);
        } else {
477 478
        if (!codec->codec_tag)
            codec->codec_tag = codec_get_tag(codec_bmp_tags, codec->codec_id);
479
        if (!codec->codec_tag) {
David Conrad's avatar
David Conrad committed
480
            av_log(s, AV_LOG_ERROR, "No bmp codec ID found.");
481 482
            ret = -1;
        }
483

484
        put_bmp_header(dyn_cp, codec, codec_bmp_tags, 0);
485
        }
486 487

    } else if (codec->codec_type == CODEC_TYPE_AUDIO) {
488
        if (!codec->codec_tag)
David Conrad's avatar
David Conrad committed
489
            codec->codec_tag = codec_get_tag(codec_wav_tags, codec->codec_id);
490
        if (!codec->codec_tag) {
David Conrad's avatar
David Conrad committed
491
            av_log(s, AV_LOG_ERROR, "No wav codec ID found.");
492
            ret = -1;
493 494
        }

495
        put_wav_header(dyn_cp, codec);
496
    }
497

498
    codecpriv_size = url_close_dyn_buf(dyn_cp, &codecpriv);
499 500 501 502
    if (codecpriv_size)
        put_ebml_binary(pb, MATROSKA_ID_CODECPRIVATE, codecpriv, codecpriv_size);
    av_free(codecpriv);
    return ret;
503 504
}

505
static int mkv_write_tracks(AVFormatContext *s)
506 507
{
    MatroskaMuxContext *mkv = s->priv_data;
508
    ByteIOContext *pb = s->pb;
509
    ebml_master tracks;
510
    int i, j, ret;
511

512 513
    ret = mkv_add_seekhead_entry(mkv->main_seekhead, MATROSKA_ID_TRACKS, url_ftell(pb));
    if (ret < 0) return ret;
514

515
    tracks = start_ebml_master(pb, MATROSKA_ID_TRACKS, 0);
516 517
    for (i = 0; i < s->nb_streams; i++) {
        AVStream *st = s->streams[i];
518
        AVCodecContext *codec = st->codec;
519
        ebml_master subinfo, track;
520
        int native_id = 0;
521
        int qt_id = 0;
522
        int bit_depth = av_get_bits_per_sample(codec->codec_id);
523 524
        int sample_rate = codec->sample_rate;
        int output_sample_rate = 0;
525
        AVMetadataTag *tag;
526

527 528 529
        if (!bit_depth)
            bit_depth = av_get_bits_per_sample_format(codec->sample_fmt);

530
        if (codec->codec_id == CODEC_ID_AAC)
531
            get_aac_sample_rates(s, codec, &sample_rate, &output_sample_rate);
532

533
        track = start_ebml_master(pb, MATROSKA_ID_TRACKENTRY, 0);
534 535
        put_ebml_uint (pb, MATROSKA_ID_TRACKNUMBER     , i + 1);
        put_ebml_uint (pb, MATROSKA_ID_TRACKUID        , i + 1);
536 537
        put_ebml_uint (pb, MATROSKA_ID_TRACKFLAGLACING , 0);    // no lacing (yet)

538 539 540 541
        if ((tag = av_metadata_get(st->metadata, "description", NULL, 0)))
            put_ebml_string(pb, MATROSKA_ID_TRACKNAME, tag->value);
        tag = av_metadata_get(st->metadata, "language", NULL, 0);
        put_ebml_string(pb, MATROSKA_ID_TRACKLANGUAGE, tag ? tag->value:"und");
542

543
        if (st->disposition)
David Conrad's avatar
David Conrad committed
544
            put_ebml_uint(pb, MATROSKA_ID_TRACKFLAGDEFAULT, !!(st->disposition & AV_DISPOSITION_DEFAULT));
545

David Conrad's avatar
David Conrad committed
546
        // look for a codec ID string specific to mkv to use,
David Conrad's avatar
David Conrad committed
547
        // if none are found, use AVI codes
548 549 550 551 552 553 554 555 556 557 558 559
        for (j = 0; ff_mkv_codec_tags[j].id != CODEC_ID_NONE; j++) {
            if (ff_mkv_codec_tags[j].id == codec->codec_id) {
                put_ebml_string(pb, MATROSKA_ID_CODECID, ff_mkv_codec_tags[j].str);
                native_id = 1;
                break;
            }
        }

        switch (codec->codec_type) {
            case CODEC_TYPE_VIDEO:
                put_ebml_uint(pb, MATROSKA_ID_TRACKTYPE, MATROSKA_TRACK_TYPE_VIDEO);

560 561 562 563 564 565 566 567 568 569 570
                if (!native_id &&
                      codec_get_tag(codec_movvideo_tags, codec->codec_id) &&
                    (!codec_get_tag(codec_bmp_tags,      codec->codec_id)
                     || codec->codec_id == CODEC_ID_SVQ1
                     || codec->codec_id == CODEC_ID_SVQ3
                     || codec->codec_id == CODEC_ID_CINEPAK))
                    qt_id = 1;

                if (qt_id)
                    put_ebml_string(pb, MATROSKA_ID_CODECID, "V_QUICKTIME");
                else if (!native_id)
David Conrad's avatar
David Conrad committed
571
                    // if there is no mkv-specific codec ID, use VFW mode
572
                    put_ebml_string(pb, MATROSKA_ID_CODECID, "V_MS/VFW/FOURCC");
David Conrad's avatar
David Conrad committed
573

574
                subinfo = start_ebml_master(pb, MATROSKA_ID_TRACKVIDEO, 0);
575 576 577
                // XXX: interlace flag?
                put_ebml_uint (pb, MATROSKA_ID_VIDEOPIXELWIDTH , codec->width);
                put_ebml_uint (pb, MATROSKA_ID_VIDEOPIXELHEIGHT, codec->height);
578 579 580 581
                if (st->sample_aspect_ratio.num) {
                    int d_width = codec->width*av_q2d(st->sample_aspect_ratio);
                    put_ebml_uint(pb, MATROSKA_ID_VIDEODISPLAYWIDTH , d_width);
                    put_ebml_uint(pb, MATROSKA_ID_VIDEODISPLAYHEIGHT, codec->height);
582
                }
583 584 585 586 587 588
                end_ebml_master(pb, subinfo);
                break;

            case CODEC_TYPE_AUDIO:
                put_ebml_uint(pb, MATROSKA_ID_TRACKTYPE, MATROSKA_TRACK_TYPE_AUDIO);

David Conrad's avatar
David Conrad committed
589
                if (!native_id)
590
                    // no mkv-specific ID, use ACM mode
591
                    put_ebml_string(pb, MATROSKA_ID_CODECID, "A_MS/ACM");
David Conrad's avatar
David Conrad committed
592

593
                subinfo = start_ebml_master(pb, MATROSKA_ID_TRACKAUDIO, 0);
594
                put_ebml_uint  (pb, MATROSKA_ID_AUDIOCHANNELS    , codec->channels);
595 596 597
                put_ebml_float (pb, MATROSKA_ID_AUDIOSAMPLINGFREQ, sample_rate);
                if (output_sample_rate)
                    put_ebml_float(pb, MATROSKA_ID_AUDIOOUTSAMPLINGFREQ, output_sample_rate);
598 599
                if (bit_depth)
                    put_ebml_uint(pb, MATROSKA_ID_AUDIOBITDEPTH, bit_depth);
600 601 602
                end_ebml_master(pb, subinfo);
                break;

David Conrad's avatar
David Conrad committed
603 604 605
            case CODEC_TYPE_SUBTITLE:
                put_ebml_uint(pb, MATROSKA_ID_TRACKTYPE, MATROSKA_TRACK_TYPE_SUBTITLE);
                break;
606
            default:
David Conrad's avatar
David Conrad committed
607
                av_log(s, AV_LOG_ERROR, "Only audio, video, and subtitles are supported for Matroska.");
608 609
                break;
        }
610
        ret = mkv_write_codecprivate(s, pb, codec, native_id, qt_id);
611 612
        if (ret < 0) return ret;

613
        end_ebml_master(pb, track);
614 615 616

        // ms precision is the de-facto standard timescale for mkv files
        av_set_pts_info(st, 64, 1, 1000);
617 618
    }
    end_ebml_master(pb, tracks);
619
    return 0;
620 621 622 623 624
}

static int mkv_write_header(AVFormatContext *s)
{
    MatroskaMuxContext *mkv = s->priv_data;
625
    ByteIOContext *pb = s->pb;
626
    ebml_master ebml_header, segment_info;
627
    AVMetadataTag *tag;
628
    int ret;
629

630 631 632
    mkv->md5_ctx = av_mallocz(av_md5_size);
    av_md5_init(mkv->md5_ctx);

633
    ebml_header = start_ebml_master(pb, EBML_ID_HEADER, 0);
634 635 636 637 638 639 640 641 642
    put_ebml_uint   (pb, EBML_ID_EBMLVERSION        ,           1);
    put_ebml_uint   (pb, EBML_ID_EBMLREADVERSION    ,           1);
    put_ebml_uint   (pb, EBML_ID_EBMLMAXIDLENGTH    ,           4);
    put_ebml_uint   (pb, EBML_ID_EBMLMAXSIZELENGTH  ,           8);
    put_ebml_string (pb, EBML_ID_DOCTYPE            ,  "matroska");
    put_ebml_uint   (pb, EBML_ID_DOCTYPEVERSION     ,           2);
    put_ebml_uint   (pb, EBML_ID_DOCTYPEREADVERSION ,           2);
    end_ebml_master(pb, ebml_header);

643
    mkv->segment = start_ebml_master(pb, MATROSKA_ID_SEGMENT, 0);
644 645
    mkv->segment_offset = url_ftell(pb);

David Conrad's avatar
David Conrad committed
646 647 648 649 650
    // we write 2 seek heads - one at the end of the file to point to each
    // cluster, and one at the beginning to point to all other level one
    // elements (including the seek head at the end of the file), which
    // isn't more than 10 elements if we only write one of each other
    // currently defined level 1 element
651 652
    mkv->main_seekhead    = mkv_start_seekhead(pb, mkv->segment_offset, 10);
    mkv->cluster_seekhead = mkv_start_seekhead(pb, mkv->segment_offset, 0);
653 654
    if (mkv->main_seekhead == NULL || mkv->cluster_seekhead == NULL)
        return AVERROR(ENOMEM);
655

656 657
    ret = mkv_add_seekhead_entry(mkv->main_seekhead, MATROSKA_ID_INFO, url_ftell(pb));
    if (ret < 0) return ret;
658

659
    segment_info = start_ebml_master(pb, MATROSKA_ID_INFO, 0);
660
    put_ebml_uint(pb, MATROSKA_ID_TIMECODESCALE, 1000000);
661 662
    if ((tag = av_metadata_get(s->metadata, "title", NULL, 0)))
        put_ebml_string(pb, MATROSKA_ID_TITLE, tag->value);
663
    if (!(s->streams[0]->codec->flags & CODEC_FLAG_BITEXACT)) {
David Conrad's avatar
David Conrad committed
664
        put_ebml_string(pb, MATROSKA_ID_MUXINGAPP , LIBAVFORMAT_IDENT);
665
        put_ebml_string(pb, MATROSKA_ID_WRITINGAPP, LIBAVFORMAT_IDENT);
666 667 668 669

        // reserve space to write the segment UID later
        mkv->segment_uid = url_ftell(pb);
        put_ebml_void(pb, 19);
670
    }
David Conrad's avatar
David Conrad committed
671

672 673 674 675 676 677
    // reserve space for the duration
    mkv->duration = 0;
    mkv->duration_offset = url_ftell(pb);
    put_ebml_void(pb, 11);                  // assumes double-precision float to be written
    end_ebml_master(pb, segment_info);

678 679
    ret = mkv_write_tracks(s);
    if (ret < 0) return ret;
680

681 682
    ret = mkv_add_seekhead_entry(mkv->cluster_seekhead, MATROSKA_ID_CLUSTER, url_ftell(pb));
    if (ret < 0) return ret;
683

David Conrad's avatar
David Conrad committed
684
    mkv->cluster_pos = url_ftell(pb);
685
    mkv->cluster = start_ebml_master(pb, MATROSKA_ID_CLUSTER, 0);
686
    put_ebml_uint(pb, MATROSKA_ID_CLUSTERTIMECODE, 0);
687
    mkv->cluster_pts = 0;
688

David Conrad's avatar
David Conrad committed
689 690
    mkv->cues = mkv_start_cues(mkv->segment_offset);
    if (mkv->cues == NULL)
691
        return AVERROR(ENOMEM);
David Conrad's avatar
David Conrad committed
692

693
    put_flush_packet(pb);
694 695 696
    return 0;
}

697
static int mkv_blockgroup_size(int pkt_size)
698
{
699
    int size = pkt_size + 4;
700
    size += ebml_num_size(size);
701 702
    size += 2;              // EBML ID for block and block duration
    size += 8;              // max size of block duration
703
    size += ebml_num_size(size);
704 705 706 707
    size += 1;              // blockgroup EBML ID
    return size;
}

708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766
static int ass_get_duration(const uint8_t *p)
{
    int sh, sm, ss, sc, eh, em, es, ec;
    uint64_t start, end;

    if (sscanf(p, "%*[^,],%d:%d:%d%*c%d,%d:%d:%d%*c%d",
               &sh, &sm, &ss, &sc, &eh, &em, &es, &ec) != 8)
        return 0;
    start = 3600000*sh + 60000*sm + 1000*ss + 10*sc;
    end   = 3600000*eh + 60000*em + 1000*es + 10*ec;
    return end - start;
}

static int mkv_write_ass_blocks(AVFormatContext *s, AVPacket *pkt)
{
    MatroskaMuxContext *mkv = s->priv_data;
    ByteIOContext *pb = s->pb;
    int i, layer = 0, max_duration = 0, size, line_size, data_size = pkt->size;
    uint8_t *start, *end, *data = pkt->data;
    ebml_master blockgroup;
    char buffer[2048];

    while (data_size) {
        int duration = ass_get_duration(data);
        max_duration = FFMAX(duration, max_duration);
        end = memchr(data, '\n', data_size);
        size = line_size = end ? end-data+1 : data_size;
        size -= end ? (end[-1]=='\r')+1 : 0;
        start = data;
        for (i=0; i<3; i++, start++)
            if (!(start = memchr(start, ',', size-(start-data))))
                return max_duration;
        size -= start - data;
        sscanf(data, "Dialogue: %d,", &layer);
        i = snprintf(buffer, sizeof(buffer), "%"PRId64",%d,",
                     s->streams[pkt->stream_index]->nb_frames++, layer);
        size = FFMIN(i+size, sizeof(buffer));
        memcpy(buffer+i, start, size-i);

        av_log(s, AV_LOG_DEBUG, "Writing block at offset %" PRIu64 ", size %d, "
               "pts %" PRId64 ", duration %d\n",
               url_ftell(pb), size, pkt->pts, duration);
        blockgroup = start_ebml_master(pb, MATROSKA_ID_BLOCKGROUP, mkv_blockgroup_size(size));
        put_ebml_id(pb, MATROSKA_ID_BLOCK);
        put_ebml_num(pb, size+4, 0);
        put_byte(pb, 0x80 | (pkt->stream_index + 1));     // this assumes stream_index is less than 126
        put_be16(pb, pkt->pts - mkv->cluster_pts);
        put_byte(pb, 0);
        put_buffer(pb, buffer, size);
        put_ebml_uint(pb, MATROSKA_ID_BLOCKDURATION, duration);
        end_ebml_master(pb, blockgroup);

        data += line_size;
        data_size -= line_size;
    }

    return max_duration;
}

767 768 769
static void mkv_write_block(AVFormatContext *s, unsigned int blockid, AVPacket *pkt, int flags)
{
    MatroskaMuxContext *mkv = s->priv_data;
770
    ByteIOContext *pb = s->pb;
771
    AVCodecContext *codec = s->streams[pkt->stream_index]->codec;
772 773
    uint8_t *data = NULL;
    int size = pkt->size;
774

775 776
    av_log(s, AV_LOG_DEBUG, "Writing block at offset %" PRIu64 ", size %d, "
           "pts %" PRId64 ", dts %" PRId64 ", duration %d, flags %d\n",
David Conrad's avatar
David Conrad committed
777
           url_ftell(pb), pkt->size, pkt->pts, pkt->dts, pkt->duration, flags);
778 779 780 781 782
    if (codec->codec_id == CODEC_ID_H264 && codec->extradata_size > 0 &&
        (AV_RB24(codec->extradata) == 1 || AV_RB32(codec->extradata) == 1))
        ff_avc_parse_nal_units_buf(pkt->data, &data, &size);
    else
        data = pkt->data;
783
    put_ebml_id(pb, blockid);
784
    put_ebml_num(pb, size+4, 0);
785 786 787
    put_byte(pb, 0x80 | (pkt->stream_index + 1));     // this assumes stream_index is less than 126
    put_be16(pb, pkt->pts - mkv->cluster_pts);
    put_byte(pb, flags);
788 789 790
    put_buffer(pb, data, size);
    if (data != pkt->data)
        av_free(data);
791 792
}

793 794
static int mkv_write_packet(AVFormatContext *s, AVPacket *pkt)
{
795
    MatroskaMuxContext *mkv = s->priv_data;
796
    ByteIOContext *pb = s->pb;
David Conrad's avatar
David Conrad committed
797
    AVCodecContext *codec = s->streams[pkt->stream_index]->codec;
David Conrad's avatar
David Conrad committed
798
    int keyframe = !!(pkt->flags & PKT_FLAG_KEY);
799
    int duration = pkt->duration;
800
    int ret;
801

802
    // start a new cluster every 5 MB or 5 sec
803
    if (url_ftell(pb) > mkv->cluster_pos + 5*1024*1024 || pkt->pts > mkv->cluster_pts + 5000) {
804 805
        av_log(s, AV_LOG_DEBUG, "Starting new cluster at offset %" PRIu64
               " bytes, pts %" PRIu64 "\n", url_ftell(pb), pkt->pts);
806
        end_ebml_master(pb, mkv->cluster);
807

808 809
        ret = mkv_add_seekhead_entry(mkv->cluster_seekhead, MATROSKA_ID_CLUSTER, url_ftell(pb));
        if (ret < 0) return ret;
810

David Conrad's avatar
David Conrad committed
811
        mkv->cluster_pos = url_ftell(pb);
812
        mkv->cluster = start_ebml_master(pb, MATROSKA_ID_CLUSTER, 0);
813 814
        put_ebml_uint(pb, MATROSKA_ID_CLUSTERTIMECODE, pkt->pts);
        mkv->cluster_pts = pkt->pts;
815
        av_md5_update(mkv->md5_ctx, pkt->data, FFMIN(200, pkt->size));
816 817
    }

David Conrad's avatar
David Conrad committed
818
    if (codec->codec_type != CODEC_TYPE_SUBTITLE) {
David Conrad's avatar
David Conrad committed
819
        mkv_write_block(s, MATROSKA_ID_SIMPLEBLOCK, pkt, keyframe << 7);
820 821
    } else if (codec->codec_id == CODEC_ID_SSA) {
        duration = mkv_write_ass_blocks(s, pkt);
David Conrad's avatar
David Conrad committed
822
    } else {
823
        ebml_master blockgroup = start_ebml_master(pb, MATROSKA_ID_BLOCKGROUP, mkv_blockgroup_size(pkt->size));
824
        duration = pkt->convergence_duration;
David Conrad's avatar
David Conrad committed
825
        mkv_write_block(s, MATROSKA_ID_BLOCK, pkt, 0);
826
        put_ebml_uint(pb, MATROSKA_ID_BLOCKDURATION, duration);
David Conrad's avatar
David Conrad committed
827 828
        end_ebml_master(pb, blockgroup);
    }
829

David Conrad's avatar
David Conrad committed
830
    if (codec->codec_type == CODEC_TYPE_VIDEO && keyframe) {
831 832
        ret = mkv_add_cuepoint(mkv->cues, pkt, mkv->cluster_pos);
        if (ret < 0) return ret;
David Conrad's avatar
David Conrad committed
833 834
    }

835
    mkv->duration = FFMAX(mkv->duration, pkt->pts + duration);
836 837 838 839 840 841
    return 0;
}

static int mkv_write_trailer(AVFormatContext *s)
{
    MatroskaMuxContext *mkv = s->priv_data;
842
    ByteIOContext *pb = s->pb;
843
    int64_t currentpos, second_seekhead, cuespos;
844
    int ret;
845

846
    end_ebml_master(pb, mkv->cluster);
847

David Conrad's avatar
David Conrad committed
848
    if (!url_is_streamed(pb)) {
David Conrad's avatar
David Conrad committed
849 850
        cuespos = mkv_write_cues(pb, mkv->cues, s->nb_streams);
        second_seekhead = mkv_write_seekhead(pb, mkv->cluster_seekhead);
David Conrad's avatar
David Conrad committed
851

David Conrad's avatar
David Conrad committed
852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871
        ret = mkv_add_seekhead_entry(mkv->main_seekhead, MATROSKA_ID_CUES    , cuespos);
        if (ret < 0) return ret;
        ret = mkv_add_seekhead_entry(mkv->main_seekhead, MATROSKA_ID_SEEKHEAD, second_seekhead);
        if (ret < 0) return ret;
        mkv_write_seekhead(pb, mkv->main_seekhead);

        // update the duration
        av_log(s, AV_LOG_DEBUG, "end duration = %" PRIu64 "\n", mkv->duration);
        currentpos = url_ftell(pb);
        url_fseek(pb, mkv->duration_offset, SEEK_SET);
        put_ebml_float(pb, MATROSKA_ID_DURATION, mkv->duration);

        // write the md5sum of some frames as the segment UID
        if (!(s->streams[0]->codec->flags & CODEC_FLAG_BITEXACT)) {
            uint8_t segment_uid[16];
            av_md5_final(mkv->md5_ctx, segment_uid);
            url_fseek(pb, mkv->segment_uid, SEEK_SET);
            put_ebml_binary(pb, MATROSKA_ID_SEGMENTUID, segment_uid, 16);
        }
        url_fseek(pb, currentpos, SEEK_SET);
David Conrad's avatar
David Conrad committed
872
    }
873

874
    end_ebml_master(pb, mkv->segment);
875
    av_free(mkv->md5_ctx);
876
    put_flush_packet(pb);
877 878 879 880 881
    return 0;
}

AVOutputFormat matroska_muxer = {
    "matroska",
882
    NULL_IF_CONFIG_SMALL("Matroska file format"),
883 884 885 886 887 888 889 890
    "video/x-matroska",
    "mkv",
    sizeof(MatroskaMuxContext),
    CODEC_ID_MP2,
    CODEC_ID_MPEG4,
    mkv_write_header,
    mkv_write_packet,
    mkv_write_trailer,
891
    .flags = AVFMT_GLOBALHEADER | AVFMT_VARIABLE_FPS,
892
    .codec_tag = (const AVCodecTag* const []){codec_bmp_tags, codec_wav_tags, 0},
893
    .subtitle_codec = CODEC_ID_TEXT,
894
};
David Conrad's avatar
David Conrad committed
895 896 897

AVOutputFormat matroska_audio_muxer = {
    "matroska",
898
    NULL_IF_CONFIG_SMALL("Matroska file format"),
David Conrad's avatar
David Conrad committed
899 900 901 902 903 904 905 906
    "audio/x-matroska",
    "mka",
    sizeof(MatroskaMuxContext),
    CODEC_ID_MP2,
    CODEC_ID_NONE,
    mkv_write_header,
    mkv_write_packet,
    mkv_write_trailer,
907
    .flags = AVFMT_GLOBALHEADER,
908
    .codec_tag = (const AVCodecTag* const []){codec_wav_tags, 0},
David Conrad's avatar
David Conrad committed
909
};