1 // Copyright (c) 2012 The WebM project authors. All Rights Reserved.
2 //
3 // Use of this source code is governed by a BSD-style license
4 // that can be found in the LICENSE file in the root of the source
5 // tree. An additional intellectual property rights grant can be found
6 // in the file PATENTS.  All contributing project authors may
7 // be found in the AUTHORS file in the root of the source tree.
8 
9 #include "mkvmuxer.hpp"
10 
11 #include <climits>
12 #include <cstdio>
13 #include <cstdlib>
14 #include <cstring>
15 #include <ctime>
16 #include <new>
17 
18 #include "mkvmuxerutil.hpp"
19 #include "mkvparser.hpp"
20 #include "mkvwriter.hpp"
21 #include "webmids.hpp"
22 
23 #ifdef _MSC_VER
24 // Disable MSVC warnings that suggest making code non-portable.
25 #pragma warning(disable:4996)
26 #endif
27 
28 namespace mkvmuxer {
29 
30 namespace {
31 // Deallocate the string designated by |dst|, and then copy the |src|
32 // string to |dst|.  The caller owns both the |src| string and the
33 // |dst| copy (hence the caller is responsible for eventually
34 // deallocating the strings, either directly, or indirectly via
35 // StrCpy).  Returns true if the source string was successfully copied
36 // to the destination.
StrCpy(const char * src,char ** dst_ptr)37 bool StrCpy(const char* src, char** dst_ptr) {
38   if (dst_ptr == NULL)
39     return false;
40 
41   char*& dst = *dst_ptr;
42 
43   delete [] dst;
44   dst = NULL;
45 
46   if (src == NULL)
47     return true;
48 
49   const size_t size = strlen(src) + 1;
50 
51   dst = new (std::nothrow) char[size];  // NOLINT
52   if (dst == NULL)
53     return false;
54 
55   strcpy(dst, src);  // NOLINT
56   return true;
57 }
58 }  // namespace
59 
60 ///////////////////////////////////////////////////////////////
61 //
62 // IMkvWriter Class
63 
IMkvWriter()64 IMkvWriter::IMkvWriter() {
65 }
66 
~IMkvWriter()67 IMkvWriter::~IMkvWriter() {
68 }
69 
WriteEbmlHeader(IMkvWriter * writer)70 bool WriteEbmlHeader(IMkvWriter* writer) {
71   // Level 0
72   uint64 size = EbmlElementSize(kMkvEBMLVersion, 1ULL);
73   size += EbmlElementSize(kMkvEBMLReadVersion, 1ULL);
74   size += EbmlElementSize(kMkvEBMLMaxIDLength, 4ULL);
75   size += EbmlElementSize(kMkvEBMLMaxSizeLength, 8ULL);
76   size += EbmlElementSize(kMkvDocType, "webm");
77   size += EbmlElementSize(kMkvDocTypeVersion, 2ULL);
78   size += EbmlElementSize(kMkvDocTypeReadVersion, 2ULL);
79 
80   if (!WriteEbmlMasterElement(writer, kMkvEBML, size))
81     return false;
82   if (!WriteEbmlElement(writer, kMkvEBMLVersion, 1ULL))
83     return false;
84   if (!WriteEbmlElement(writer, kMkvEBMLReadVersion, 1ULL))
85     return false;
86   if (!WriteEbmlElement(writer, kMkvEBMLMaxIDLength, 4ULL))
87     return false;
88   if (!WriteEbmlElement(writer, kMkvEBMLMaxSizeLength, 8ULL))
89     return false;
90   if (!WriteEbmlElement(writer, kMkvDocType, "webm"))
91     return false;
92   if (!WriteEbmlElement(writer, kMkvDocTypeVersion, 2ULL))
93     return false;
94   if (!WriteEbmlElement(writer, kMkvDocTypeReadVersion, 2ULL))
95     return false;
96 
97   return true;
98 }
99 
ChunkedCopy(mkvparser::IMkvReader * source,mkvmuxer::IMkvWriter * dst,mkvmuxer::int64 start,int64 size)100 bool ChunkedCopy(mkvparser::IMkvReader* source,
101                  mkvmuxer::IMkvWriter* dst,
102                  mkvmuxer::int64 start, int64 size) {
103   // TODO(vigneshv): Check if this is a reasonable value.
104   const uint32 kBufSize = 2048;
105   uint8* buf = new uint8[kBufSize];
106   int64 offset = start;
107   while (size > 0) {
108     const int64 read_len = (size > kBufSize) ? kBufSize : size;
109     if (source->Read(offset, static_cast<long>(read_len), buf))
110       return false;
111     dst->Write(buf, static_cast<uint32>(read_len));
112     offset += read_len;
113     size -= read_len;
114   }
115   delete[] buf;
116   return true;
117 }
118 
119 ///////////////////////////////////////////////////////////////
120 //
121 // Frame Class
122 
Frame()123 Frame::Frame()
124     : add_id_(0),
125       additional_(NULL),
126       additional_length_(0),
127       duration_(0),
128       frame_(NULL),
129       is_key_(false),
130       length_(0),
131       track_number_(0),
132       timestamp_(0),
133       discard_padding_(0) {
134 }
135 
~Frame()136 Frame::~Frame() {
137   delete [] frame_;
138   delete [] additional_;
139 }
140 
Init(const uint8 * frame,uint64 length)141 bool Frame::Init(const uint8* frame, uint64 length) {
142   uint8* const data =
143       new (std::nothrow) uint8[static_cast<size_t>(length)];  // NOLINT
144   if (!data)
145     return false;
146 
147   delete [] frame_;
148   frame_ = data;
149   length_ = length;
150 
151   memcpy(frame_, frame, static_cast<size_t>(length_));
152   return true;
153 }
154 
AddAdditionalData(const uint8 * additional,uint64 length,uint64 add_id)155 bool Frame::AddAdditionalData(const uint8* additional, uint64 length,
156                               uint64 add_id) {
157   uint8* const data =
158       new (std::nothrow) uint8[static_cast<size_t>(length)];  // NOLINT
159   if (!data)
160     return false;
161 
162   delete [] additional_;
163   additional_ = data;
164   additional_length_ = length;
165   add_id_ = add_id;
166 
167   memcpy(additional_, additional, static_cast<size_t>(additional_length_));
168   return true;
169 }
170 
171 ///////////////////////////////////////////////////////////////
172 //
173 // CuePoint Class
174 
CuePoint()175 CuePoint::CuePoint()
176     : time_(0),
177       track_(0),
178       cluster_pos_(0),
179       block_number_(1),
180       output_block_number_(true) {
181 }
182 
~CuePoint()183 CuePoint::~CuePoint() {
184 }
185 
Write(IMkvWriter * writer) const186 bool CuePoint::Write(IMkvWriter* writer) const {
187   if (!writer || track_ < 1 || cluster_pos_ < 1)
188     return false;
189 
190   uint64 size = EbmlElementSize(kMkvCueClusterPosition, cluster_pos_);
191   size += EbmlElementSize(kMkvCueTrack, track_);
192   if (output_block_number_ && block_number_ > 1)
193     size += EbmlElementSize(kMkvCueBlockNumber, block_number_);
194   const uint64 track_pos_size = EbmlMasterElementSize(kMkvCueTrackPositions,
195                                                       size) + size;
196   const uint64 payload_size = EbmlElementSize(kMkvCueTime, time_) +
197                               track_pos_size;
198 
199   if (!WriteEbmlMasterElement(writer, kMkvCuePoint, payload_size))
200     return false;
201 
202   const int64 payload_position = writer->Position();
203   if (payload_position < 0)
204     return false;
205 
206   if (!WriteEbmlElement(writer, kMkvCueTime, time_))
207     return false;
208 
209   if (!WriteEbmlMasterElement(writer, kMkvCueTrackPositions, size))
210     return false;
211   if (!WriteEbmlElement(writer, kMkvCueTrack, track_))
212     return false;
213   if (!WriteEbmlElement(writer, kMkvCueClusterPosition, cluster_pos_))
214     return false;
215   if (output_block_number_ && block_number_ > 1)
216     if (!WriteEbmlElement(writer, kMkvCueBlockNumber, block_number_))
217       return false;
218 
219   const int64 stop_position = writer->Position();
220   if (stop_position < 0)
221     return false;
222 
223   if (stop_position - payload_position != static_cast<int64>(payload_size))
224     return false;
225 
226   return true;
227 }
228 
PayloadSize() const229 uint64 CuePoint::PayloadSize() const {
230   uint64 size = EbmlElementSize(kMkvCueClusterPosition, cluster_pos_);
231   size += EbmlElementSize(kMkvCueTrack, track_);
232   if (output_block_number_ && block_number_ > 1)
233     size += EbmlElementSize(kMkvCueBlockNumber, block_number_);
234   const uint64 track_pos_size = EbmlMasterElementSize(kMkvCueTrackPositions,
235                                                       size) + size;
236   const uint64 payload_size = EbmlElementSize(kMkvCueTime, time_) +
237                               track_pos_size;
238 
239   return payload_size;
240 }
241 
Size() const242 uint64 CuePoint::Size() const {
243   const uint64 payload_size = PayloadSize();
244   return EbmlMasterElementSize(kMkvCuePoint, payload_size) + payload_size;
245 }
246 
247 ///////////////////////////////////////////////////////////////
248 //
249 // Cues Class
250 
Cues()251 Cues::Cues()
252     : cue_entries_capacity_(0),
253       cue_entries_size_(0),
254       cue_entries_(NULL),
255       output_block_number_(true) {
256 }
257 
~Cues()258 Cues::~Cues() {
259   if (cue_entries_) {
260     for (int32 i = 0; i < cue_entries_size_; ++i) {
261       CuePoint* const cue = cue_entries_[i];
262       delete cue;
263     }
264     delete [] cue_entries_;
265   }
266 }
267 
AddCue(CuePoint * cue)268 bool Cues::AddCue(CuePoint* cue) {
269   if (!cue)
270     return false;
271 
272   if ((cue_entries_size_ + 1) > cue_entries_capacity_) {
273     // Add more CuePoints.
274     const int32 new_capacity =
275         (!cue_entries_capacity_) ? 2 : cue_entries_capacity_ * 2;
276 
277     if (new_capacity < 1)
278       return false;
279 
280     CuePoint** const cues =
281         new (std::nothrow) CuePoint*[new_capacity];  // NOLINT
282     if (!cues)
283       return false;
284 
285     for (int32 i = 0; i < cue_entries_size_; ++i) {
286       cues[i] = cue_entries_[i];
287     }
288 
289     delete [] cue_entries_;
290 
291     cue_entries_ = cues;
292     cue_entries_capacity_ = new_capacity;
293   }
294 
295   cue->set_output_block_number(output_block_number_);
296   cue_entries_[cue_entries_size_++] = cue;
297   return true;
298 }
299 
GetCueByIndex(int32 index) const300 CuePoint* Cues::GetCueByIndex(int32 index) const {
301   if (cue_entries_ == NULL)
302     return NULL;
303 
304   if (index >= cue_entries_size_)
305     return NULL;
306 
307   return cue_entries_[index];
308 }
309 
Size()310 uint64 Cues::Size() {
311   uint64 size = 0;
312   for (int32 i = 0; i < cue_entries_size_; ++i)
313     size += GetCueByIndex(i)->Size();
314   size += EbmlMasterElementSize(kMkvCues, size);
315   return size;
316 }
317 
Write(IMkvWriter * writer) const318 bool Cues::Write(IMkvWriter* writer) const {
319   if (!writer)
320     return false;
321 
322   uint64 size = 0;
323   for (int32 i = 0; i < cue_entries_size_; ++i) {
324     const CuePoint* const cue = GetCueByIndex(i);
325 
326     if (!cue)
327       return false;
328 
329     size += cue->Size();
330   }
331 
332   if (!WriteEbmlMasterElement(writer, kMkvCues, size))
333     return false;
334 
335   const int64 payload_position = writer->Position();
336   if (payload_position < 0)
337     return false;
338 
339   for (int32 i = 0; i < cue_entries_size_; ++i) {
340     const CuePoint* const cue = GetCueByIndex(i);
341 
342     if (!cue->Write(writer))
343       return false;
344   }
345 
346   const int64 stop_position = writer->Position();
347   if (stop_position < 0)
348     return false;
349 
350   if (stop_position - payload_position != static_cast<int64>(size))
351     return false;
352 
353   return true;
354 }
355 
356 ///////////////////////////////////////////////////////////////
357 //
358 // ContentEncAESSettings Class
359 
ContentEncAESSettings()360 ContentEncAESSettings::ContentEncAESSettings() : cipher_mode_(kCTR) {}
361 
Size() const362 uint64 ContentEncAESSettings::Size() const {
363   const uint64 payload = PayloadSize();
364   const uint64 size =
365       EbmlMasterElementSize(kMkvContentEncAESSettings, payload) + payload;
366   return size;
367 }
368 
Write(IMkvWriter * writer) const369 bool ContentEncAESSettings::Write(IMkvWriter* writer) const {
370   const uint64 payload = PayloadSize();
371 
372   if (!WriteEbmlMasterElement(writer, kMkvContentEncAESSettings, payload))
373     return false;
374 
375   const int64 payload_position = writer->Position();
376   if (payload_position < 0)
377     return false;
378 
379   if (!WriteEbmlElement(writer, kMkvAESSettingsCipherMode, cipher_mode_))
380     return false;
381 
382   const int64 stop_position = writer->Position();
383   if (stop_position < 0 ||
384       stop_position - payload_position != static_cast<int64>(payload))
385     return false;
386 
387   return true;
388 }
389 
PayloadSize() const390 uint64 ContentEncAESSettings::PayloadSize() const {
391   uint64 size = EbmlElementSize(kMkvAESSettingsCipherMode, cipher_mode_);
392   return size;
393 }
394 
395 ///////////////////////////////////////////////////////////////
396 //
397 // ContentEncoding Class
398 
ContentEncoding()399 ContentEncoding::ContentEncoding()
400     : enc_algo_(5),
401       enc_key_id_(NULL),
402       encoding_order_(0),
403       encoding_scope_(1),
404       encoding_type_(1),
405       enc_key_id_length_(0) {
406 }
407 
~ContentEncoding()408 ContentEncoding::~ContentEncoding() {
409   delete [] enc_key_id_;
410 }
411 
SetEncryptionID(const uint8 * id,uint64 length)412 bool ContentEncoding::SetEncryptionID(const uint8* id, uint64 length) {
413   if (!id || length < 1)
414     return false;
415 
416   delete [] enc_key_id_;
417 
418   enc_key_id_ =
419       new (std::nothrow) uint8[static_cast<size_t>(length)];  // NOLINT
420   if (!enc_key_id_)
421     return false;
422 
423   memcpy(enc_key_id_, id, static_cast<size_t>(length));
424   enc_key_id_length_ = length;
425 
426   return true;
427 }
428 
Size() const429 uint64 ContentEncoding::Size() const {
430   const uint64 encryption_size = EncryptionSize();
431   const uint64 encoding_size = EncodingSize(0, encryption_size);
432   const uint64 encodings_size = EbmlMasterElementSize(kMkvContentEncoding,
433                                                       encoding_size) +
434                                 encoding_size;
435 
436   return encodings_size;
437 }
438 
Write(IMkvWriter * writer) const439 bool ContentEncoding::Write(IMkvWriter* writer) const {
440   const uint64 encryption_size = EncryptionSize();
441   const uint64 encoding_size = EncodingSize(0, encryption_size);
442   const uint64 size = EbmlMasterElementSize(kMkvContentEncoding,
443                                             encoding_size) +
444                       encoding_size;
445 
446   const int64 payload_position = writer->Position();
447   if (payload_position < 0)
448     return false;
449 
450   if (!WriteEbmlMasterElement(writer, kMkvContentEncoding, encoding_size))
451     return false;
452   if (!WriteEbmlElement(writer, kMkvContentEncodingOrder, encoding_order_))
453     return false;
454   if (!WriteEbmlElement(writer, kMkvContentEncodingScope, encoding_scope_))
455     return false;
456   if (!WriteEbmlElement(writer, kMkvContentEncodingType, encoding_type_))
457     return false;
458 
459   if (!WriteEbmlMasterElement(writer, kMkvContentEncryption, encryption_size))
460     return false;
461   if (!WriteEbmlElement(writer, kMkvContentEncAlgo, enc_algo_))
462     return false;
463   if (!WriteEbmlElement(writer,
464                         kMkvContentEncKeyID,
465                         enc_key_id_,
466                         enc_key_id_length_))
467     return false;
468 
469   if (!enc_aes_settings_.Write(writer))
470     return false;
471 
472   const int64 stop_position = writer->Position();
473   if (stop_position < 0 ||
474       stop_position - payload_position != static_cast<int64>(size))
475     return false;
476 
477   return true;
478 }
479 
EncodingSize(uint64 compresion_size,uint64 encryption_size) const480 uint64 ContentEncoding::EncodingSize(uint64 compresion_size,
481                                      uint64 encryption_size) const {
482   // TODO(fgalligan): Add support for compression settings.
483   if (compresion_size != 0)
484     return 0;
485 
486   uint64 encoding_size = 0;
487 
488   if (encryption_size > 0) {
489     encoding_size += EbmlMasterElementSize(kMkvContentEncryption,
490                                            encryption_size) +
491                      encryption_size;
492   }
493   encoding_size += EbmlElementSize(kMkvContentEncodingType, encoding_type_);
494   encoding_size += EbmlElementSize(kMkvContentEncodingScope, encoding_scope_);
495   encoding_size += EbmlElementSize(kMkvContentEncodingOrder, encoding_order_);
496 
497   return encoding_size;
498 }
499 
EncryptionSize() const500 uint64 ContentEncoding::EncryptionSize() const {
501   const uint64 aes_size = enc_aes_settings_.Size();
502 
503   uint64 encryption_size = EbmlElementSize(kMkvContentEncKeyID,
504                                            enc_key_id_,
505                                            enc_key_id_length_);
506   encryption_size += EbmlElementSize(kMkvContentEncAlgo, enc_algo_);
507 
508   return encryption_size + aes_size;
509 }
510 
511 ///////////////////////////////////////////////////////////////
512 //
513 // Track Class
514 
Track(unsigned int * seed)515 Track::Track(unsigned int* seed)
516     : codec_id_(NULL),
517       codec_private_(NULL),
518       language_(NULL),
519       max_block_additional_id_(0),
520       name_(NULL),
521       number_(0),
522       type_(0),
523       uid_(MakeUID(seed)),
524       codec_delay_(0),
525       seek_pre_roll_(0),
526       codec_private_length_(0),
527       content_encoding_entries_(NULL),
528       content_encoding_entries_size_(0) {
529 }
530 
~Track()531 Track::~Track() {
532   delete [] codec_id_;
533   delete [] codec_private_;
534   delete [] language_;
535   delete [] name_;
536 
537   if (content_encoding_entries_) {
538     for (uint32 i = 0; i < content_encoding_entries_size_; ++i) {
539       ContentEncoding* const encoding = content_encoding_entries_[i];
540       delete encoding;
541     }
542     delete [] content_encoding_entries_;
543   }
544 }
545 
AddContentEncoding()546 bool Track::AddContentEncoding() {
547   const uint32 count = content_encoding_entries_size_ + 1;
548 
549   ContentEncoding** const content_encoding_entries =
550       new (std::nothrow) ContentEncoding*[count];  // NOLINT
551   if (!content_encoding_entries)
552     return false;
553 
554   ContentEncoding* const content_encoding =
555       new (std::nothrow) ContentEncoding();  // NOLINT
556   if (!content_encoding) {
557     delete [] content_encoding_entries;
558     return false;
559   }
560 
561   for (uint32 i = 0; i < content_encoding_entries_size_; ++i) {
562     content_encoding_entries[i] = content_encoding_entries_[i];
563   }
564 
565   delete [] content_encoding_entries_;
566 
567   content_encoding_entries_ = content_encoding_entries;
568   content_encoding_entries_[content_encoding_entries_size_] = content_encoding;
569   content_encoding_entries_size_ = count;
570   return true;
571 }
572 
GetContentEncodingByIndex(uint32 index) const573 ContentEncoding* Track::GetContentEncodingByIndex(uint32 index) const {
574   if (content_encoding_entries_ == NULL)
575     return NULL;
576 
577   if (index >= content_encoding_entries_size_)
578     return NULL;
579 
580   return content_encoding_entries_[index];
581 }
582 
PayloadSize() const583 uint64 Track::PayloadSize() const {
584   uint64 size = EbmlElementSize(kMkvTrackNumber, number_);
585   size += EbmlElementSize(kMkvTrackUID, uid_);
586   size += EbmlElementSize(kMkvTrackType, type_);
587   if (codec_id_)
588     size += EbmlElementSize(kMkvCodecID, codec_id_);
589   if (codec_private_)
590     size += EbmlElementSize(kMkvCodecPrivate,
591                             codec_private_,
592                             codec_private_length_);
593   if (language_)
594     size += EbmlElementSize(kMkvLanguage, language_);
595   if (name_)
596     size += EbmlElementSize(kMkvName, name_);
597   if (max_block_additional_id_)
598     size += EbmlElementSize(kMkvMaxBlockAdditionID, max_block_additional_id_);
599   if (codec_delay_)
600     size += EbmlElementSize(kMkvCodecDelay, codec_delay_);
601   if (seek_pre_roll_)
602     size += EbmlElementSize(kMkvSeekPreRoll, seek_pre_roll_);
603 
604   if (content_encoding_entries_size_ > 0) {
605     uint64 content_encodings_size = 0;
606     for (uint32 i = 0; i < content_encoding_entries_size_; ++i) {
607       ContentEncoding* const encoding = content_encoding_entries_[i];
608       content_encodings_size += encoding->Size();
609     }
610 
611     size += EbmlMasterElementSize(kMkvContentEncodings,
612                                   content_encodings_size) +
613             content_encodings_size;
614   }
615 
616   return size;
617 }
618 
Size() const619 uint64 Track::Size() const {
620   uint64 size = PayloadSize();
621   size += EbmlMasterElementSize(kMkvTrackEntry, size);
622   return size;
623 }
624 
Write(IMkvWriter * writer) const625 bool Track::Write(IMkvWriter* writer) const {
626   if (!writer)
627     return false;
628 
629   // |size| may be bigger than what is written out in this function because
630   // derived classes may write out more data in the Track element.
631   const uint64 payload_size = PayloadSize();
632 
633   if (!WriteEbmlMasterElement(writer, kMkvTrackEntry, payload_size))
634     return false;
635 
636   uint64 size = EbmlElementSize(kMkvTrackNumber, number_);
637   size += EbmlElementSize(kMkvTrackUID, uid_);
638   size += EbmlElementSize(kMkvTrackType, type_);
639   if (codec_id_)
640     size += EbmlElementSize(kMkvCodecID, codec_id_);
641   if (codec_private_)
642     size += EbmlElementSize(kMkvCodecPrivate,
643                             codec_private_,
644                             codec_private_length_);
645   if (language_)
646     size += EbmlElementSize(kMkvLanguage, language_);
647   if (name_)
648     size += EbmlElementSize(kMkvName, name_);
649   if (max_block_additional_id_)
650     size += EbmlElementSize(kMkvMaxBlockAdditionID, max_block_additional_id_);
651   if (codec_delay_)
652     size += EbmlElementSize(kMkvCodecDelay, codec_delay_);
653   if (seek_pre_roll_)
654     size += EbmlElementSize(kMkvSeekPreRoll, seek_pre_roll_);
655 
656 
657   const int64 payload_position = writer->Position();
658   if (payload_position < 0)
659     return false;
660 
661   if (!WriteEbmlElement(writer, kMkvTrackNumber, number_))
662     return false;
663   if (!WriteEbmlElement(writer, kMkvTrackUID, uid_))
664     return false;
665   if (!WriteEbmlElement(writer, kMkvTrackType, type_))
666     return false;
667   if (max_block_additional_id_) {
668     if (!WriteEbmlElement(writer,
669                           kMkvMaxBlockAdditionID,
670                           max_block_additional_id_)) {
671       return false;
672     }
673   }
674   if (codec_delay_) {
675     if (!WriteEbmlElement(writer, kMkvCodecDelay, codec_delay_))
676       return false;
677   }
678   if (seek_pre_roll_) {
679     if (!WriteEbmlElement(writer, kMkvSeekPreRoll, seek_pre_roll_))
680       return false;
681   }
682   if (codec_id_) {
683     if (!WriteEbmlElement(writer, kMkvCodecID, codec_id_))
684       return false;
685   }
686   if (codec_private_) {
687     if (!WriteEbmlElement(writer,
688                           kMkvCodecPrivate,
689                           codec_private_,
690                           codec_private_length_))
691       return false;
692   }
693   if (language_) {
694     if (!WriteEbmlElement(writer, kMkvLanguage, language_))
695       return false;
696   }
697   if (name_) {
698     if (!WriteEbmlElement(writer, kMkvName, name_))
699       return false;
700   }
701 
702   int64 stop_position = writer->Position();
703   if (stop_position < 0 ||
704       stop_position - payload_position != static_cast<int64>(size))
705     return false;
706 
707   if (content_encoding_entries_size_ > 0) {
708     uint64 content_encodings_size = 0;
709     for (uint32 i = 0; i < content_encoding_entries_size_; ++i) {
710       ContentEncoding* const encoding = content_encoding_entries_[i];
711       content_encodings_size += encoding->Size();
712     }
713 
714     if (!WriteEbmlMasterElement(writer,
715                                 kMkvContentEncodings,
716                                 content_encodings_size))
717       return false;
718 
719     for (uint32 i = 0; i < content_encoding_entries_size_; ++i) {
720       ContentEncoding* const encoding = content_encoding_entries_[i];
721       if (!encoding->Write(writer))
722         return false;
723     }
724   }
725 
726   stop_position = writer->Position();
727   if (stop_position < 0)
728     return false;
729   return true;
730 }
731 
SetCodecPrivate(const uint8 * codec_private,uint64 length)732 bool Track::SetCodecPrivate(const uint8* codec_private, uint64 length) {
733   if (!codec_private || length < 1)
734     return false;
735 
736   delete [] codec_private_;
737 
738   codec_private_ =
739       new (std::nothrow) uint8[static_cast<size_t>(length)];  // NOLINT
740   if (!codec_private_)
741     return false;
742 
743   memcpy(codec_private_, codec_private, static_cast<size_t>(length));
744   codec_private_length_ = length;
745 
746   return true;
747 }
748 
set_codec_id(const char * codec_id)749 void Track::set_codec_id(const char* codec_id) {
750   if (codec_id) {
751     delete [] codec_id_;
752 
753     const size_t length = strlen(codec_id) + 1;
754     codec_id_ = new (std::nothrow) char[length];  // NOLINT
755     if (codec_id_) {
756 #ifdef _MSC_VER
757       strcpy_s(codec_id_, length, codec_id);
758 #else
759       strcpy(codec_id_, codec_id);
760 #endif
761     }
762   }
763 }
764 
765 // TODO(fgalligan): Vet the language parameter.
set_language(const char * language)766 void Track::set_language(const char* language) {
767   if (language) {
768     delete [] language_;
769 
770     const size_t length = strlen(language) + 1;
771     language_ = new (std::nothrow) char[length];  // NOLINT
772     if (language_) {
773 #ifdef _MSC_VER
774       strcpy_s(language_, length, language);
775 #else
776       strcpy(language_, language);
777 #endif
778     }
779   }
780 }
781 
set_name(const char * name)782 void Track::set_name(const char* name) {
783   if (name) {
784     delete [] name_;
785 
786     const size_t length = strlen(name) + 1;
787     name_ = new (std::nothrow) char[length];  // NOLINT
788     if (name_) {
789 #ifdef _MSC_VER
790       strcpy_s(name_, length, name);
791 #else
792       strcpy(name_, name);
793 #endif
794     }
795   }
796 }
797 
798 ///////////////////////////////////////////////////////////////
799 //
800 // VideoTrack Class
801 
VideoTrack(unsigned int * seed)802 VideoTrack::VideoTrack(unsigned int* seed)
803     : Track(seed),
804       display_height_(0),
805       display_width_(0),
806       frame_rate_(0.0),
807       height_(0),
808       stereo_mode_(0),
809       alpha_mode_(0),
810       width_(0) {
811 }
812 
~VideoTrack()813 VideoTrack::~VideoTrack() {
814 }
815 
SetStereoMode(uint64 stereo_mode)816 bool VideoTrack::SetStereoMode(uint64 stereo_mode) {
817   if (stereo_mode != kMono &&
818       stereo_mode != kSideBySideLeftIsFirst &&
819       stereo_mode != kTopBottomRightIsFirst &&
820       stereo_mode != kTopBottomLeftIsFirst &&
821       stereo_mode != kSideBySideRightIsFirst)
822     return false;
823 
824   stereo_mode_ = stereo_mode;
825   return true;
826 }
827 
SetAlphaMode(uint64 alpha_mode)828 bool VideoTrack::SetAlphaMode(uint64 alpha_mode) {
829   if (alpha_mode != kNoAlpha &&
830       alpha_mode != kAlpha)
831     return false;
832 
833   alpha_mode_ = alpha_mode;
834   return true;
835 }
836 
PayloadSize() const837 uint64 VideoTrack::PayloadSize() const {
838   const uint64 parent_size = Track::PayloadSize();
839 
840   uint64 size = VideoPayloadSize();
841   size += EbmlMasterElementSize(kMkvVideo, size);
842 
843   return parent_size + size;
844 }
845 
Write(IMkvWriter * writer) const846 bool VideoTrack::Write(IMkvWriter* writer) const {
847   if (!Track::Write(writer))
848     return false;
849 
850   const uint64 size = VideoPayloadSize();
851 
852   if (!WriteEbmlMasterElement(writer, kMkvVideo, size))
853     return false;
854 
855   const int64 payload_position = writer->Position();
856   if (payload_position < 0)
857     return false;
858 
859   if (!WriteEbmlElement(writer, kMkvPixelWidth, width_))
860     return false;
861   if (!WriteEbmlElement(writer, kMkvPixelHeight, height_))
862     return false;
863   if (display_width_ > 0)
864     if (!WriteEbmlElement(writer, kMkvDisplayWidth, display_width_))
865       return false;
866   if (display_height_ > 0)
867     if (!WriteEbmlElement(writer, kMkvDisplayHeight, display_height_))
868       return false;
869   if (stereo_mode_ > kMono)
870     if (!WriteEbmlElement(writer, kMkvStereoMode, stereo_mode_))
871       return false;
872   if (alpha_mode_ > kNoAlpha)
873     if (!WriteEbmlElement(writer, kMkvAlphaMode, alpha_mode_))
874       return false;
875   if (frame_rate_ > 0.0)
876     if (!WriteEbmlElement(writer,
877                           kMkvFrameRate,
878                           static_cast<float>(frame_rate_)))
879       return false;
880 
881   const int64 stop_position = writer->Position();
882   if (stop_position < 0 ||
883       stop_position - payload_position != static_cast<int64>(size))
884     return false;
885 
886   return true;
887 }
888 
VideoPayloadSize() const889 uint64 VideoTrack::VideoPayloadSize() const {
890   uint64 size = EbmlElementSize(kMkvPixelWidth, width_);
891   size += EbmlElementSize(kMkvPixelHeight, height_);
892   if (display_width_ > 0)
893     size += EbmlElementSize(kMkvDisplayWidth, display_width_);
894   if (display_height_ > 0)
895     size += EbmlElementSize(kMkvDisplayHeight, display_height_);
896   if (stereo_mode_ > kMono)
897     size += EbmlElementSize(kMkvStereoMode, stereo_mode_);
898   if (alpha_mode_ > kNoAlpha)
899     size += EbmlElementSize(kMkvAlphaMode, alpha_mode_);
900   if (frame_rate_ > 0.0)
901     size += EbmlElementSize(kMkvFrameRate, static_cast<float>(frame_rate_));
902 
903   return size;
904 }
905 
906 ///////////////////////////////////////////////////////////////
907 //
908 // AudioTrack Class
909 
AudioTrack(unsigned int * seed)910 AudioTrack::AudioTrack(unsigned int* seed)
911     : Track(seed),
912       bit_depth_(0),
913       channels_(1),
914       sample_rate_(0.0) {
915 }
916 
~AudioTrack()917 AudioTrack::~AudioTrack() {
918 }
919 
PayloadSize() const920 uint64 AudioTrack::PayloadSize() const {
921   const uint64 parent_size = Track::PayloadSize();
922 
923   uint64 size = EbmlElementSize(kMkvSamplingFrequency,
924                                 static_cast<float>(sample_rate_));
925   size += EbmlElementSize(kMkvChannels, channels_);
926   if (bit_depth_ > 0)
927     size += EbmlElementSize(kMkvBitDepth, bit_depth_);
928   size += EbmlMasterElementSize(kMkvAudio, size);
929 
930   return parent_size + size;
931 }
932 
Write(IMkvWriter * writer) const933 bool AudioTrack::Write(IMkvWriter* writer) const {
934   if (!Track::Write(writer))
935     return false;
936 
937   // Calculate AudioSettings size.
938   uint64 size = EbmlElementSize(kMkvSamplingFrequency,
939                                 static_cast<float>(sample_rate_));
940   size += EbmlElementSize(kMkvChannels, channels_);
941   if (bit_depth_ > 0)
942     size += EbmlElementSize(kMkvBitDepth, bit_depth_);
943 
944   if (!WriteEbmlMasterElement(writer, kMkvAudio, size))
945     return false;
946 
947   const int64 payload_position = writer->Position();
948   if (payload_position < 0)
949     return false;
950 
951   if (!WriteEbmlElement(writer,
952                         kMkvSamplingFrequency,
953                         static_cast<float>(sample_rate_)))
954     return false;
955   if (!WriteEbmlElement(writer, kMkvChannels, channels_))
956     return false;
957   if (bit_depth_ > 0)
958     if (!WriteEbmlElement(writer, kMkvBitDepth, bit_depth_))
959       return false;
960 
961   const int64 stop_position = writer->Position();
962   if (stop_position < 0 ||
963       stop_position - payload_position != static_cast<int64>(size))
964     return false;
965 
966   return true;
967 }
968 
969 ///////////////////////////////////////////////////////////////
970 //
971 // Tracks Class
972 
973 const char Tracks::kOpusCodecId[] = "A_OPUS";
974 const char Tracks::kVorbisCodecId[] = "A_VORBIS";
975 const char Tracks::kVp8CodecId[] = "V_VP8";
976 const char Tracks::kVp9CodecId[] = "V_VP9";
977 
978 
Tracks()979 Tracks::Tracks()
980     : track_entries_(NULL),
981       track_entries_size_(0) {
982 }
983 
~Tracks()984 Tracks::~Tracks() {
985   if (track_entries_) {
986     for (uint32 i = 0; i < track_entries_size_; ++i) {
987       Track* const track = track_entries_[i];
988       delete track;
989     }
990     delete [] track_entries_;
991   }
992 }
993 
AddTrack(Track * track,int32 number)994 bool Tracks::AddTrack(Track* track, int32 number) {
995   if (number < 0)
996     return false;
997 
998   // This muxer only supports track numbers in the range [1, 126], in
999   // order to be able (to use Matroska integer representation) to
1000   // serialize the block header (of which the track number is a part)
1001   // for a frame using exactly 4 bytes.
1002 
1003   if (number > 0x7E)
1004     return false;
1005 
1006   uint32 track_num = number;
1007 
1008   if (track_num > 0) {
1009     // Check to make sure a track does not already have |track_num|.
1010     for (uint32 i = 0; i < track_entries_size_; ++i) {
1011       if (track_entries_[i]->number() == track_num)
1012         return false;
1013     }
1014   }
1015 
1016   const uint32 count = track_entries_size_ + 1;
1017 
1018   Track** const track_entries = new (std::nothrow) Track*[count];  // NOLINT
1019   if (!track_entries)
1020     return false;
1021 
1022   for (uint32 i = 0; i < track_entries_size_; ++i) {
1023     track_entries[i] = track_entries_[i];
1024   }
1025 
1026   delete [] track_entries_;
1027 
1028   // Find the lowest availible track number > 0.
1029   if (track_num == 0) {
1030     track_num = count;
1031 
1032     // Check to make sure a track does not already have |track_num|.
1033     bool exit = false;
1034     do {
1035       exit = true;
1036       for (uint32 i = 0; i < track_entries_size_; ++i) {
1037         if (track_entries[i]->number() == track_num) {
1038           track_num++;
1039           exit = false;
1040           break;
1041         }
1042       }
1043     } while (!exit);
1044   }
1045   track->set_number(track_num);
1046 
1047   track_entries_ = track_entries;
1048   track_entries_[track_entries_size_] = track;
1049   track_entries_size_ = count;
1050   return true;
1051 }
1052 
GetTrackByIndex(uint32 index) const1053 const Track* Tracks::GetTrackByIndex(uint32 index) const {
1054   if (track_entries_ == NULL)
1055     return NULL;
1056 
1057   if (index >= track_entries_size_)
1058     return NULL;
1059 
1060   return track_entries_[index];
1061 }
1062 
GetTrackByNumber(uint64 track_number) const1063 Track* Tracks::GetTrackByNumber(uint64 track_number) const {
1064   const int32 count = track_entries_size();
1065   for (int32 i = 0; i < count; ++i) {
1066     if (track_entries_[i]->number() == track_number)
1067       return track_entries_[i];
1068   }
1069 
1070   return NULL;
1071 }
1072 
TrackIsAudio(uint64 track_number) const1073 bool Tracks::TrackIsAudio(uint64 track_number) const {
1074   const Track* const track = GetTrackByNumber(track_number);
1075 
1076   if (track->type() == kAudio)
1077     return true;
1078 
1079   return false;
1080 }
1081 
TrackIsVideo(uint64 track_number) const1082 bool Tracks::TrackIsVideo(uint64 track_number) const {
1083   const Track* const track = GetTrackByNumber(track_number);
1084 
1085   if (track->type() == kVideo)
1086     return true;
1087 
1088   return false;
1089 }
1090 
Write(IMkvWriter * writer) const1091 bool Tracks::Write(IMkvWriter* writer) const {
1092   uint64 size = 0;
1093   const int32 count = track_entries_size();
1094   for (int32 i = 0; i < count; ++i) {
1095     const Track* const track = GetTrackByIndex(i);
1096 
1097     if (!track)
1098       return false;
1099 
1100     size += track->Size();
1101   }
1102 
1103   if (!WriteEbmlMasterElement(writer, kMkvTracks, size))
1104     return false;
1105 
1106   const int64 payload_position = writer->Position();
1107   if (payload_position < 0)
1108     return false;
1109 
1110   for (int32 i = 0; i < count; ++i) {
1111     const Track* const track = GetTrackByIndex(i);
1112     if (!track->Write(writer))
1113       return false;
1114   }
1115 
1116   const int64 stop_position = writer->Position();
1117   if (stop_position < 0 ||
1118       stop_position - payload_position != static_cast<int64>(size))
1119     return false;
1120 
1121   return true;
1122 }
1123 
1124 ///////////////////////////////////////////////////////////////
1125 //
1126 // Chapter Class
1127 
set_id(const char * id)1128 bool Chapter::set_id(const char* id) {
1129   return StrCpy(id, &id_);
1130 }
1131 
set_time(const Segment & segment,uint64 start_ns,uint64 end_ns)1132 void Chapter::set_time(const Segment& segment,
1133                        uint64 start_ns,
1134                        uint64 end_ns) {
1135   const SegmentInfo* const info = segment.GetSegmentInfo();
1136   const uint64 timecode_scale = info->timecode_scale();
1137   start_timecode_ = start_ns / timecode_scale;
1138   end_timecode_ = end_ns / timecode_scale;
1139 }
1140 
add_string(const char * title,const char * language,const char * country)1141 bool Chapter::add_string(const char* title,
1142                          const char* language,
1143                          const char* country) {
1144   if (!ExpandDisplaysArray())
1145     return false;
1146 
1147   Display& d = displays_[displays_count_++];
1148   d.Init();
1149 
1150   if (!d.set_title(title))
1151     return false;
1152 
1153   if (!d.set_language(language))
1154     return false;
1155 
1156   if (!d.set_country(country))
1157     return false;
1158 
1159   return true;
1160 }
1161 
Chapter()1162 Chapter::Chapter() {
1163   // This ctor only constructs the object.  Proper initialization is
1164   // done in Init() (called in Chapters::AddChapter()).  The only
1165   // reason we bother implementing this ctor is because we had to
1166   // declare it as private (along with the dtor), in order to prevent
1167   // clients from creating Chapter instances (a privelege we grant
1168   // only to the Chapters class).  Doing no initialization here also
1169   // means that creating arrays of chapter objects is more efficient,
1170   // because we only initialize each new chapter object as it becomes
1171   // active on the array.
1172 }
1173 
~Chapter()1174 Chapter::~Chapter() {
1175 }
1176 
Init(unsigned int * seed)1177 void Chapter::Init(unsigned int* seed) {
1178   id_ = NULL;
1179   displays_ = NULL;
1180   displays_size_ = 0;
1181   displays_count_ = 0;
1182   uid_ = MakeUID(seed);
1183 }
1184 
ShallowCopy(Chapter * dst) const1185 void Chapter::ShallowCopy(Chapter* dst) const {
1186   dst->id_ = id_;
1187   dst->start_timecode_ = start_timecode_;
1188   dst->end_timecode_ = end_timecode_;
1189   dst->uid_ = uid_;
1190   dst->displays_ = displays_;
1191   dst->displays_size_ = displays_size_;
1192   dst->displays_count_ = displays_count_;
1193 }
1194 
Clear()1195 void Chapter::Clear() {
1196   StrCpy(NULL, &id_);
1197 
1198   while (displays_count_ > 0) {
1199     Display& d = displays_[--displays_count_];
1200     d.Clear();
1201   }
1202 
1203   delete [] displays_;
1204   displays_ = NULL;
1205 
1206   displays_size_ = 0;
1207 }
1208 
ExpandDisplaysArray()1209 bool Chapter::ExpandDisplaysArray() {
1210   if (displays_size_ > displays_count_)
1211     return true;  // nothing to do yet
1212 
1213   const int size = (displays_size_ == 0) ? 1 : 2 * displays_size_;
1214 
1215   Display* const displays = new (std::nothrow) Display[size];  // NOLINT
1216   if (displays == NULL)
1217     return false;
1218 
1219   for (int idx = 0; idx < displays_count_; ++idx) {
1220     displays[idx] = displays_[idx];  // shallow copy
1221   }
1222 
1223   delete [] displays_;
1224 
1225   displays_ = displays;
1226   displays_size_ = size;
1227 
1228   return true;
1229 }
1230 
WriteAtom(IMkvWriter * writer) const1231 uint64 Chapter::WriteAtom(IMkvWriter* writer) const {
1232   uint64 payload_size =
1233       EbmlElementSize(kMkvChapterStringUID, id_) +
1234       EbmlElementSize(kMkvChapterUID, uid_) +
1235       EbmlElementSize(kMkvChapterTimeStart, start_timecode_) +
1236       EbmlElementSize(kMkvChapterTimeEnd, end_timecode_);
1237 
1238   for (int idx = 0; idx < displays_count_; ++idx) {
1239     const Display& d = displays_[idx];
1240     payload_size += d.WriteDisplay(NULL);
1241   }
1242 
1243   const uint64 atom_size =
1244       EbmlMasterElementSize(kMkvChapterAtom, payload_size) +
1245       payload_size;
1246 
1247   if (writer == NULL)
1248     return atom_size;
1249 
1250   const int64 start = writer->Position();
1251 
1252   if (!WriteEbmlMasterElement(writer, kMkvChapterAtom, payload_size))
1253     return 0;
1254 
1255   if (!WriteEbmlElement(writer, kMkvChapterStringUID, id_))
1256     return 0;
1257 
1258   if (!WriteEbmlElement(writer, kMkvChapterUID, uid_))
1259     return 0;
1260 
1261   if (!WriteEbmlElement(writer, kMkvChapterTimeStart, start_timecode_))
1262     return 0;
1263 
1264   if (!WriteEbmlElement(writer, kMkvChapterTimeEnd, end_timecode_))
1265     return 0;
1266 
1267   for (int idx = 0; idx < displays_count_; ++idx) {
1268     const Display& d = displays_[idx];
1269 
1270     if (!d.WriteDisplay(writer))
1271       return 0;
1272   }
1273 
1274   const int64 stop = writer->Position();
1275 
1276   if (stop >= start && uint64(stop - start) != atom_size)
1277     return 0;
1278 
1279   return atom_size;
1280 }
1281 
Init()1282 void Chapter::Display::Init() {
1283   title_ = NULL;
1284   language_ = NULL;
1285   country_ = NULL;
1286 }
1287 
Clear()1288 void Chapter::Display::Clear() {
1289   StrCpy(NULL, &title_);
1290   StrCpy(NULL, &language_);
1291   StrCpy(NULL, &country_);
1292 }
1293 
set_title(const char * title)1294 bool Chapter::Display::set_title(const char* title) {
1295   return StrCpy(title, &title_);
1296 }
1297 
set_language(const char * language)1298 bool Chapter::Display::set_language(const char* language) {
1299   return StrCpy(language, &language_);
1300 }
1301 
set_country(const char * country)1302 bool Chapter::Display::set_country(const char* country) {
1303   return StrCpy(country, &country_);
1304 }
1305 
WriteDisplay(IMkvWriter * writer) const1306 uint64 Chapter::Display::WriteDisplay(IMkvWriter* writer) const {
1307   uint64 payload_size = EbmlElementSize(kMkvChapString, title_);
1308 
1309   if (language_)
1310     payload_size += EbmlElementSize(kMkvChapLanguage, language_);
1311 
1312   if (country_)
1313     payload_size += EbmlElementSize(kMkvChapCountry, country_);
1314 
1315   const uint64 display_size =
1316       EbmlMasterElementSize(kMkvChapterDisplay, payload_size) +
1317       payload_size;
1318 
1319   if (writer == NULL)
1320     return display_size;
1321 
1322   const int64 start = writer->Position();
1323 
1324   if (!WriteEbmlMasterElement(writer, kMkvChapterDisplay, payload_size))
1325     return 0;
1326 
1327   if (!WriteEbmlElement(writer, kMkvChapString, title_))
1328     return 0;
1329 
1330   if (language_) {
1331     if (!WriteEbmlElement(writer, kMkvChapLanguage, language_))
1332       return 0;
1333   }
1334 
1335   if (country_) {
1336     if (!WriteEbmlElement(writer, kMkvChapCountry, country_))
1337       return 0;
1338   }
1339 
1340   const int64 stop = writer->Position();
1341 
1342   if (stop >= start && uint64(stop - start) != display_size)
1343     return 0;
1344 
1345   return display_size;
1346 }
1347 
1348 ///////////////////////////////////////////////////////////////
1349 //
1350 // Chapters Class
1351 
Chapters()1352 Chapters::Chapters()
1353     : chapters_size_(0),
1354       chapters_count_(0),
1355       chapters_(NULL) {
1356 }
1357 
~Chapters()1358 Chapters::~Chapters() {
1359   while (chapters_count_ > 0) {
1360     Chapter& chapter = chapters_[--chapters_count_];
1361     chapter.Clear();
1362   }
1363 
1364   delete [] chapters_;
1365   chapters_ = NULL;
1366 }
1367 
Count() const1368 int Chapters::Count() const {
1369   return chapters_count_;
1370 }
1371 
AddChapter(unsigned int * seed)1372 Chapter* Chapters::AddChapter(unsigned int* seed) {
1373   if (!ExpandChaptersArray())
1374     return NULL;
1375 
1376   Chapter& chapter = chapters_[chapters_count_++];
1377   chapter.Init(seed);
1378 
1379   return &chapter;
1380 }
1381 
Write(IMkvWriter * writer) const1382 bool Chapters::Write(IMkvWriter* writer) const {
1383   if (writer == NULL)
1384     return false;
1385 
1386   const uint64 payload_size = WriteEdition(NULL);  // return size only
1387 
1388   if (!WriteEbmlMasterElement(writer, kMkvChapters, payload_size))
1389     return false;
1390 
1391   const int64 start = writer->Position();
1392 
1393   if (WriteEdition(writer) == 0)  // error
1394     return false;
1395 
1396   const int64 stop = writer->Position();
1397 
1398   if (stop >= start && uint64(stop - start) != payload_size)
1399     return false;
1400 
1401   return true;
1402 }
1403 
ExpandChaptersArray()1404 bool Chapters::ExpandChaptersArray() {
1405   if (chapters_size_ > chapters_count_)
1406     return true;  // nothing to do yet
1407 
1408   const int size = (chapters_size_ == 0) ? 1 : 2 * chapters_size_;
1409 
1410   Chapter* const chapters = new (std::nothrow) Chapter[size];  // NOLINT
1411   if (chapters == NULL)
1412     return false;
1413 
1414   for (int idx = 0; idx < chapters_count_; ++idx) {
1415     const Chapter& src = chapters_[idx];
1416     Chapter* const dst = chapters + idx;
1417     src.ShallowCopy(dst);
1418   }
1419 
1420   delete [] chapters_;
1421 
1422   chapters_ = chapters;
1423   chapters_size_ = size;
1424 
1425   return true;
1426 }
1427 
WriteEdition(IMkvWriter * writer) const1428 uint64 Chapters::WriteEdition(IMkvWriter* writer) const {
1429   uint64 payload_size = 0;
1430 
1431   for (int idx = 0; idx < chapters_count_; ++idx) {
1432     const Chapter& chapter = chapters_[idx];
1433     payload_size += chapter.WriteAtom(NULL);
1434   }
1435 
1436   const uint64 edition_size =
1437       EbmlMasterElementSize(kMkvEditionEntry, payload_size) +
1438       payload_size;
1439 
1440   if (writer == NULL)  // return size only
1441     return edition_size;
1442 
1443   const int64 start = writer->Position();
1444 
1445   if (!WriteEbmlMasterElement(writer, kMkvEditionEntry, payload_size))
1446     return 0;  // error
1447 
1448   for (int idx = 0; idx < chapters_count_; ++idx) {
1449     const Chapter& chapter = chapters_[idx];
1450 
1451     const uint64 chapter_size = chapter.WriteAtom(writer);
1452     if (chapter_size == 0)  // error
1453       return 0;
1454   }
1455 
1456   const int64 stop = writer->Position();
1457 
1458   if (stop >= start && uint64(stop - start) != edition_size)
1459     return 0;
1460 
1461   return edition_size;
1462 }
1463 
1464 ///////////////////////////////////////////////////////////////
1465 //
1466 // Cluster class
1467 
Cluster(uint64 timecode,int64 cues_pos)1468 Cluster::Cluster(uint64 timecode, int64 cues_pos)
1469     : blocks_added_(0),
1470       finalized_(false),
1471       header_written_(false),
1472       payload_size_(0),
1473       position_for_cues_(cues_pos),
1474       size_position_(-1),
1475       timecode_(timecode),
1476       writer_(NULL) {
1477 }
1478 
~Cluster()1479 Cluster::~Cluster() {
1480 }
1481 
Init(IMkvWriter * ptr_writer)1482 bool Cluster::Init(IMkvWriter* ptr_writer) {
1483   if (!ptr_writer) {
1484     return false;
1485   }
1486   writer_ = ptr_writer;
1487   return true;
1488 }
1489 
AddFrame(const uint8 * frame,uint64 length,uint64 track_number,uint64 abs_timecode,bool is_key)1490 bool Cluster::AddFrame(const uint8* frame,
1491                        uint64 length,
1492                        uint64 track_number,
1493                        uint64 abs_timecode,
1494                        bool is_key) {
1495   return DoWriteBlock(frame,
1496                       length,
1497                       track_number,
1498                       abs_timecode,
1499                       is_key ? 1 : 0,
1500                       &WriteSimpleBlock);
1501 }
1502 
AddFrameWithAdditional(const uint8 * frame,uint64 length,const uint8 * additional,uint64 additional_length,uint64 add_id,uint64 track_number,uint64 abs_timecode,bool is_key)1503 bool Cluster::AddFrameWithAdditional(const uint8* frame,
1504                                      uint64 length,
1505                                      const uint8* additional,
1506                                      uint64 additional_length,
1507                                      uint64 add_id,
1508                                      uint64 track_number,
1509                                      uint64 abs_timecode,
1510                                      bool is_key) {
1511   return DoWriteBlockWithAdditional(frame,
1512                                     length,
1513                                     additional,
1514                                     additional_length,
1515                                     add_id,
1516                                     track_number,
1517                                     abs_timecode,
1518                                     is_key ? 1 : 0,
1519                                     &WriteBlockWithAdditional);
1520 }
1521 
AddFrameWithDiscardPadding(const uint8 * frame,uint64 length,int64 discard_padding,uint64 track_number,uint64 abs_timecode,bool is_key)1522 bool Cluster::AddFrameWithDiscardPadding(const uint8* frame,
1523                                          uint64 length,
1524                                          int64 discard_padding,
1525                                          uint64 track_number,
1526                                          uint64 abs_timecode,
1527                                          bool is_key) {
1528   return DoWriteBlockWithDiscardPadding(frame,
1529                                         length,
1530                                         discard_padding,
1531                                         track_number,
1532                                         abs_timecode,
1533                                         is_key ? 1 : 0,
1534                                         &WriteBlockWithDiscardPadding);
1535 }
1536 
AddMetadata(const uint8 * frame,uint64 length,uint64 track_number,uint64 abs_timecode,uint64 duration_timecode)1537 bool Cluster::AddMetadata(const uint8* frame,
1538                           uint64 length,
1539                           uint64 track_number,
1540                           uint64 abs_timecode,
1541                           uint64 duration_timecode) {
1542   return DoWriteBlock(frame,
1543                       length,
1544                       track_number,
1545                       abs_timecode,
1546                       duration_timecode,
1547                       &WriteMetadataBlock);
1548 }
1549 
AddPayloadSize(uint64 size)1550 void Cluster::AddPayloadSize(uint64 size) {
1551   payload_size_ += size;
1552 }
1553 
Finalize()1554 bool Cluster::Finalize() {
1555   if (!writer_ || finalized_ || size_position_ == -1)
1556     return false;
1557 
1558   if (writer_->Seekable()) {
1559     const int64 pos = writer_->Position();
1560 
1561     if (writer_->Position(size_position_))
1562       return false;
1563 
1564     if (WriteUIntSize(writer_, payload_size(), 8))
1565       return false;
1566 
1567     if (writer_->Position(pos))
1568       return false;
1569   }
1570 
1571   finalized_ = true;
1572 
1573   return true;
1574 }
1575 
Size() const1576 uint64 Cluster::Size() const {
1577   const uint64 element_size =
1578       EbmlMasterElementSize(kMkvCluster,
1579                             0xFFFFFFFFFFFFFFFFULL) + payload_size_;
1580   return element_size;
1581 }
1582 
1583 template <typename Type>
PreWriteBlock(Type * write_function)1584 bool Cluster::PreWriteBlock(Type* write_function) {
1585   if (write_function == NULL)
1586     return false;
1587 
1588   if (finalized_)
1589     return false;
1590 
1591   if (!header_written_) {
1592     if (!WriteClusterHeader())
1593       return false;
1594   }
1595 
1596   return true;
1597 }
1598 
PostWriteBlock(uint64 element_size)1599 void Cluster::PostWriteBlock(uint64 element_size) {
1600   AddPayloadSize(element_size);
1601   ++blocks_added_;
1602 }
1603 
IsValidTrackNumber(uint64 track_number) const1604 bool Cluster::IsValidTrackNumber(uint64 track_number) const {
1605   return (track_number > 0 && track_number <= 0x7E);
1606 }
1607 
GetRelativeTimecode(int64 abs_timecode) const1608 int64 Cluster::GetRelativeTimecode(int64 abs_timecode) const {
1609   const int64 cluster_timecode = this->Cluster::timecode();
1610   const int64 rel_timecode =
1611       static_cast<int64>(abs_timecode) - cluster_timecode;
1612 
1613   if (rel_timecode < 0 || rel_timecode > kMaxBlockTimecode)
1614     return -1;
1615 
1616   return rel_timecode;
1617 }
1618 
DoWriteBlock(const uint8 * frame,uint64 length,uint64 track_number,uint64 abs_timecode,uint64 generic_arg,WriteBlock write_block)1619 bool Cluster::DoWriteBlock(
1620     const uint8* frame,
1621     uint64 length,
1622     uint64 track_number,
1623     uint64 abs_timecode,
1624     uint64 generic_arg,
1625     WriteBlock write_block) {
1626   if (frame == NULL || length == 0)
1627     return false;
1628 
1629   if (!IsValidTrackNumber(track_number))
1630     return false;
1631 
1632   const int64 rel_timecode = GetRelativeTimecode(abs_timecode);
1633   if (rel_timecode < 0)
1634     return false;
1635 
1636   if (!PreWriteBlock(write_block))
1637     return false;
1638 
1639   const uint64 element_size = (*write_block)(writer_,
1640                                              frame,
1641                                              length,
1642                                              track_number,
1643                                              rel_timecode,
1644                                              generic_arg);
1645   if (element_size == 0)
1646     return false;
1647 
1648   PostWriteBlock(element_size);
1649   return true;
1650 }
1651 
DoWriteBlockWithAdditional(const uint8 * frame,uint64 length,const uint8 * additional,uint64 additional_length,uint64 add_id,uint64 track_number,uint64 abs_timecode,uint64 generic_arg,WriteBlockAdditional write_block)1652 bool Cluster::DoWriteBlockWithAdditional(
1653     const uint8* frame,
1654     uint64 length,
1655     const uint8* additional,
1656     uint64 additional_length,
1657     uint64 add_id,
1658     uint64 track_number,
1659     uint64 abs_timecode,
1660     uint64 generic_arg,
1661     WriteBlockAdditional write_block) {
1662   if (frame == NULL || length == 0 ||
1663       additional == NULL || additional_length == 0)
1664     return false;
1665 
1666   if (!IsValidTrackNumber(track_number))
1667     return false;
1668 
1669   const int64 rel_timecode = GetRelativeTimecode(abs_timecode);
1670   if (rel_timecode < 0)
1671     return false;
1672 
1673   if (!PreWriteBlock(write_block))
1674     return false;
1675 
1676   const uint64 element_size = (*write_block)(writer_,
1677                                              frame,
1678                                              length,
1679                                              additional,
1680                                              additional_length,
1681                                              add_id,
1682                                              track_number,
1683                                              rel_timecode,
1684                                              generic_arg);
1685   if (element_size == 0)
1686     return false;
1687 
1688   PostWriteBlock(element_size);
1689   return true;
1690 }
1691 
DoWriteBlockWithDiscardPadding(const uint8 * frame,uint64 length,int64 discard_padding,uint64 track_number,uint64 abs_timecode,uint64 generic_arg,WriteBlockDiscardPadding write_block)1692 bool Cluster::DoWriteBlockWithDiscardPadding(
1693     const uint8* frame,
1694     uint64 length,
1695     int64 discard_padding,
1696     uint64 track_number,
1697     uint64 abs_timecode,
1698     uint64 generic_arg,
1699     WriteBlockDiscardPadding write_block) {
1700   if (frame == NULL || length == 0 || discard_padding <= 0)
1701     return false;
1702 
1703   if (!IsValidTrackNumber(track_number))
1704     return false;
1705 
1706   const int64 rel_timecode = GetRelativeTimecode(abs_timecode);
1707   if (rel_timecode < 0)
1708     return false;
1709 
1710   if (!PreWriteBlock(write_block))
1711     return false;
1712 
1713   const uint64 element_size = (*write_block)(writer_,
1714                                              frame,
1715                                              length,
1716                                              discard_padding,
1717                                              track_number,
1718                                              rel_timecode,
1719                                              generic_arg);
1720   if (element_size == 0)
1721     return false;
1722 
1723   PostWriteBlock(element_size);
1724   return true;
1725 }
1726 
WriteClusterHeader()1727 bool Cluster::WriteClusterHeader() {
1728   if (finalized_)
1729     return false;
1730 
1731   if (WriteID(writer_, kMkvCluster))
1732     return false;
1733 
1734   // Save for later.
1735   size_position_ = writer_->Position();
1736 
1737   // Write "unknown" (EBML coded -1) as cluster size value. We need to write 8
1738   // bytes because we do not know how big our cluster will be.
1739   if (SerializeInt(writer_, kEbmlUnknownValue, 8))
1740     return false;
1741 
1742   if (!WriteEbmlElement(writer_, kMkvTimecode, timecode()))
1743     return false;
1744   AddPayloadSize(EbmlElementSize(kMkvTimecode, timecode()));
1745   header_written_ = true;
1746 
1747   return true;
1748 }
1749 
1750 ///////////////////////////////////////////////////////////////
1751 //
1752 // SeekHead Class
1753 
SeekHead()1754 SeekHead::SeekHead() : start_pos_(0ULL) {
1755   for (int32 i = 0; i < kSeekEntryCount; ++i) {
1756     seek_entry_id_[i] = 0;
1757     seek_entry_pos_[i] = 0;
1758   }
1759 }
1760 
~SeekHead()1761 SeekHead::~SeekHead() {
1762 }
1763 
Finalize(IMkvWriter * writer) const1764 bool SeekHead::Finalize(IMkvWriter* writer) const {
1765   if (writer->Seekable()) {
1766     if (start_pos_ == -1)
1767       return false;
1768 
1769     uint64 payload_size = 0;
1770     uint64 entry_size[kSeekEntryCount];
1771 
1772     for (int32 i = 0; i < kSeekEntryCount; ++i) {
1773       if (seek_entry_id_[i] != 0) {
1774         entry_size[i] = EbmlElementSize(
1775             kMkvSeekID,
1776             static_cast<uint64>(seek_entry_id_[i]));
1777         entry_size[i] += EbmlElementSize(kMkvSeekPosition, seek_entry_pos_[i]);
1778 
1779         payload_size += EbmlMasterElementSize(kMkvSeek, entry_size[i]) +
1780                         entry_size[i];
1781       }
1782     }
1783 
1784     // No SeekHead elements
1785     if (payload_size == 0)
1786       return true;
1787 
1788     const int64 pos = writer->Position();
1789     if (writer->Position(start_pos_))
1790       return false;
1791 
1792     if (!WriteEbmlMasterElement(writer, kMkvSeekHead, payload_size))
1793       return false;
1794 
1795     for (int32 i = 0; i < kSeekEntryCount; ++i) {
1796       if (seek_entry_id_[i] != 0) {
1797         if (!WriteEbmlMasterElement(writer, kMkvSeek, entry_size[i]))
1798           return false;
1799 
1800         if (!WriteEbmlElement(writer,
1801                               kMkvSeekID,
1802                               static_cast<uint64>(seek_entry_id_[i])))
1803           return false;
1804 
1805         if (!WriteEbmlElement(writer, kMkvSeekPosition, seek_entry_pos_[i]))
1806           return false;
1807       }
1808     }
1809 
1810     const uint64 total_entry_size = kSeekEntryCount * MaxEntrySize();
1811     const uint64 total_size =
1812         EbmlMasterElementSize(kMkvSeekHead,
1813                               total_entry_size) + total_entry_size;
1814     const int64 size_left = total_size - (writer->Position() - start_pos_);
1815 
1816     const uint64 bytes_written = WriteVoidElement(writer, size_left);
1817     if (!bytes_written)
1818       return false;
1819 
1820     if (writer->Position(pos))
1821       return false;
1822   }
1823 
1824   return true;
1825 }
1826 
Write(IMkvWriter * writer)1827 bool SeekHead::Write(IMkvWriter* writer) {
1828   const uint64 entry_size = kSeekEntryCount * MaxEntrySize();
1829   const uint64 size = EbmlMasterElementSize(kMkvSeekHead, entry_size);
1830 
1831   start_pos_ = writer->Position();
1832 
1833   const uint64 bytes_written = WriteVoidElement(writer, size + entry_size);
1834   if (!bytes_written)
1835     return false;
1836 
1837   return true;
1838 }
1839 
AddSeekEntry(uint32 id,uint64 pos)1840 bool SeekHead::AddSeekEntry(uint32 id, uint64 pos) {
1841   for (int32 i = 0; i < kSeekEntryCount; ++i) {
1842     if (seek_entry_id_[i] == 0) {
1843       seek_entry_id_[i] = id;
1844       seek_entry_pos_[i] = pos;
1845       return true;
1846     }
1847   }
1848   return false;
1849 }
1850 
GetId(int index) const1851 uint32 SeekHead::GetId(int index) const {
1852   if (index < 0 || index >= kSeekEntryCount)
1853     return UINT_MAX;
1854   return seek_entry_id_[index];
1855 }
1856 
GetPosition(int index) const1857 uint64 SeekHead::GetPosition(int index) const {
1858   if (index < 0 || index >= kSeekEntryCount)
1859     return ULLONG_MAX;
1860   return seek_entry_pos_[index];
1861 }
1862 
SetSeekEntry(int index,uint32 id,uint64 position)1863 bool SeekHead::SetSeekEntry(int index, uint32 id, uint64 position) {
1864   if (index < 0 || index >= kSeekEntryCount)
1865     return false;
1866   seek_entry_id_[index] = id;
1867   seek_entry_pos_[index] = position;
1868   return true;
1869 }
1870 
MaxEntrySize() const1871 uint64 SeekHead::MaxEntrySize() const {
1872   const uint64 max_entry_payload_size =
1873       EbmlElementSize(kMkvSeekID, 0xffffffffULL) +
1874       EbmlElementSize(kMkvSeekPosition, 0xffffffffffffffffULL);
1875   const uint64 max_entry_size =
1876       EbmlMasterElementSize(kMkvSeek, max_entry_payload_size) +
1877       max_entry_payload_size;
1878 
1879   return max_entry_size;
1880 }
1881 
1882 ///////////////////////////////////////////////////////////////
1883 //
1884 // SegmentInfo Class
1885 
SegmentInfo()1886 SegmentInfo::SegmentInfo()
1887     : duration_(-1.0),
1888       muxing_app_(NULL),
1889       timecode_scale_(1000000ULL),
1890       writing_app_(NULL),
1891       duration_pos_(-1) {
1892 }
1893 
~SegmentInfo()1894 SegmentInfo::~SegmentInfo() {
1895   delete [] muxing_app_;
1896   delete [] writing_app_;
1897 }
1898 
Init()1899 bool SegmentInfo::Init() {
1900   int32 major;
1901   int32 minor;
1902   int32 build;
1903   int32 revision;
1904   GetVersion(&major, &minor, &build, &revision);
1905   char temp[256];
1906 #ifdef _MSC_VER
1907   sprintf_s(temp,
1908             sizeof(temp)/sizeof(temp[0]),
1909             "libwebm-%d.%d.%d.%d",
1910             major,
1911             minor,
1912             build,
1913             revision);
1914 #else
1915   snprintf(temp,
1916            sizeof(temp)/sizeof(temp[0]),
1917            "libwebm-%d.%d.%d.%d",
1918            major,
1919            minor,
1920            build,
1921            revision);
1922 #endif
1923 
1924   const size_t app_len = strlen(temp) + 1;
1925 
1926   delete [] muxing_app_;
1927 
1928   muxing_app_ = new (std::nothrow) char[app_len];  // NOLINT
1929   if (!muxing_app_)
1930     return false;
1931 
1932 #ifdef _MSC_VER
1933   strcpy_s(muxing_app_, app_len, temp);
1934 #else
1935   strcpy(muxing_app_, temp);
1936 #endif
1937 
1938   set_writing_app(temp);
1939   if (!writing_app_)
1940     return false;
1941   return true;
1942 }
1943 
Finalize(IMkvWriter * writer) const1944 bool SegmentInfo::Finalize(IMkvWriter* writer) const {
1945   if (!writer)
1946     return false;
1947 
1948   if (duration_ > 0.0) {
1949     if (writer->Seekable()) {
1950       if (duration_pos_ == -1)
1951         return false;
1952 
1953       const int64 pos = writer->Position();
1954 
1955       if (writer->Position(duration_pos_))
1956         return false;
1957 
1958       if (!WriteEbmlElement(writer,
1959                             kMkvDuration,
1960                             static_cast<float>(duration_)))
1961         return false;
1962 
1963       if (writer->Position(pos))
1964         return false;
1965     }
1966   }
1967 
1968   return true;
1969 }
1970 
Write(IMkvWriter * writer)1971 bool SegmentInfo::Write(IMkvWriter* writer) {
1972   if (!writer || !muxing_app_ || !writing_app_)
1973     return false;
1974 
1975   uint64 size = EbmlElementSize(kMkvTimecodeScale, timecode_scale_);
1976   if (duration_ > 0.0)
1977     size += EbmlElementSize(kMkvDuration, static_cast<float>(duration_));
1978   size += EbmlElementSize(kMkvMuxingApp, muxing_app_);
1979   size += EbmlElementSize(kMkvWritingApp, writing_app_);
1980 
1981   if (!WriteEbmlMasterElement(writer, kMkvInfo, size))
1982     return false;
1983 
1984   const int64 payload_position = writer->Position();
1985   if (payload_position < 0)
1986     return false;
1987 
1988   if (!WriteEbmlElement(writer, kMkvTimecodeScale, timecode_scale_))
1989     return false;
1990 
1991   if (duration_ > 0.0) {
1992     // Save for later
1993     duration_pos_ = writer->Position();
1994 
1995     if (!WriteEbmlElement(writer, kMkvDuration, static_cast<float>(duration_)))
1996       return false;
1997   }
1998 
1999   if (!WriteEbmlElement(writer, kMkvMuxingApp, muxing_app_))
2000     return false;
2001   if (!WriteEbmlElement(writer, kMkvWritingApp, writing_app_))
2002     return false;
2003 
2004   const int64 stop_position = writer->Position();
2005   if (stop_position < 0 ||
2006       stop_position - payload_position != static_cast<int64>(size))
2007     return false;
2008 
2009   return true;
2010 }
2011 
set_muxing_app(const char * app)2012 void SegmentInfo::set_muxing_app(const char* app) {
2013   if (app) {
2014     const size_t length = strlen(app) + 1;
2015     char* temp_str = new (std::nothrow) char[length];  // NOLINT
2016     if (!temp_str)
2017       return;
2018 
2019 #ifdef _MSC_VER
2020     strcpy_s(temp_str, length, app);
2021 #else
2022     strcpy(temp_str, app);
2023 #endif
2024 
2025     delete [] muxing_app_;
2026     muxing_app_ = temp_str;
2027   }
2028 }
2029 
set_writing_app(const char * app)2030 void SegmentInfo::set_writing_app(const char* app) {
2031   if (app) {
2032     const size_t length = strlen(app) + 1;
2033     char* temp_str = new (std::nothrow) char[length];  // NOLINT
2034     if (!temp_str)
2035       return;
2036 
2037 #ifdef _MSC_VER
2038     strcpy_s(temp_str, length, app);
2039 #else
2040     strcpy(temp_str, app);
2041 #endif
2042 
2043     delete [] writing_app_;
2044     writing_app_ = temp_str;
2045   }
2046 }
2047 
2048 ///////////////////////////////////////////////////////////////
2049 //
2050 // Segment Class
2051 
Segment()2052 Segment::Segment()
2053     : chunk_count_(0),
2054       chunk_name_(NULL),
2055       chunk_writer_cluster_(NULL),
2056       chunk_writer_cues_(NULL),
2057       chunk_writer_header_(NULL),
2058       chunking_(false),
2059       chunking_base_name_(NULL),
2060       cluster_list_(NULL),
2061       cluster_list_capacity_(0),
2062       cluster_list_size_(0),
2063       cues_position_(kAfterClusters),
2064       cues_track_(0),
2065       force_new_cluster_(false),
2066       frames_(NULL),
2067       frames_capacity_(0),
2068       frames_size_(0),
2069       has_video_(false),
2070       header_written_(false),
2071       last_block_duration_(0),
2072       last_timestamp_(0),
2073       max_cluster_duration_(kDefaultMaxClusterDuration),
2074       max_cluster_size_(0),
2075       mode_(kFile),
2076       new_cuepoint_(false),
2077       output_cues_(true),
2078       payload_pos_(0),
2079       size_position_(0),
2080       writer_cluster_(NULL),
2081       writer_cues_(NULL),
2082       writer_header_(NULL) {
2083   const time_t curr_time = time(NULL);
2084   seed_ = static_cast<unsigned int>(curr_time);
2085 #ifdef _WIN32
2086   srand(seed_);
2087 #endif
2088 }
2089 
~Segment()2090 Segment::~Segment() {
2091   if (cluster_list_) {
2092     for (int32 i = 0; i < cluster_list_size_; ++i) {
2093       Cluster* const cluster = cluster_list_[i];
2094       delete cluster;
2095     }
2096     delete [] cluster_list_;
2097   }
2098 
2099   if (frames_) {
2100     for (int32 i = 0; i < frames_size_; ++i) {
2101       Frame* const frame = frames_[i];
2102       delete frame;
2103     }
2104     delete [] frames_;
2105   }
2106 
2107   delete [] chunk_name_;
2108   delete [] chunking_base_name_;
2109 
2110   if (chunk_writer_cluster_) {
2111     chunk_writer_cluster_->Close();
2112     delete chunk_writer_cluster_;
2113   }
2114   if (chunk_writer_cues_) {
2115     chunk_writer_cues_->Close();
2116     delete chunk_writer_cues_;
2117   }
2118   if (chunk_writer_header_) {
2119     chunk_writer_header_->Close();
2120     delete chunk_writer_header_;
2121   }
2122 }
2123 
MoveCuesBeforeClustersHelper(uint64 diff,int32 index,uint64 * cues_size)2124 void Segment::MoveCuesBeforeClustersHelper(uint64 diff,
2125                                            int32 index,
2126                                            uint64* cues_size) {
2127   const uint64 old_cues_size = *cues_size;
2128   CuePoint* const cue_point = cues_.GetCueByIndex(index);
2129   if (cue_point == NULL)
2130     return;
2131   const uint64 old_cue_point_size = cue_point->Size();
2132   const uint64 cluster_pos = cue_point->cluster_pos() + diff;
2133   cue_point->set_cluster_pos(cluster_pos);  // update the new cluster position
2134   // New size of the cue is computed as follows
2135   //    Let a = current size of Cues Element
2136   //    Let b = Difference in Cue Point's size after this pass
2137   //    Let c = Difference in length of Cues Element's size
2138   //            (This is computed as CodedSize(a + b) - CodedSize(a)
2139   //    Let d = a + b + c. Now d is the new size of the Cues element which is
2140   //                       passed on to the next recursive call.
2141   const uint64 cue_point_size_diff = cue_point->Size() - old_cue_point_size;
2142   const uint64 cue_size_diff = GetCodedUIntSize(*cues_size +
2143                                                 cue_point_size_diff) -
2144                                GetCodedUIntSize(*cues_size);
2145   *cues_size += cue_point_size_diff + cue_size_diff;
2146   diff = *cues_size - old_cues_size;
2147   if (diff > 0) {
2148     for (int32 i = 0; i < cues_.cue_entries_size(); ++i) {
2149       MoveCuesBeforeClustersHelper(diff, i, cues_size);
2150     }
2151   }
2152 }
2153 
MoveCuesBeforeClusters()2154 void Segment::MoveCuesBeforeClusters() {
2155   const uint64 current_cue_size = cues_.Size();
2156   uint64 cue_size = current_cue_size;
2157   for (int32 i = 0; i < cues_.cue_entries_size(); i++)
2158     MoveCuesBeforeClustersHelper(current_cue_size, i, &cue_size);
2159 
2160   // Adjust the Seek Entry to reflect the change in position
2161   // of Cluster and Cues
2162   int32 cluster_index = 0;
2163   int32 cues_index = 0;
2164   for (int32 i = 0; i < SeekHead::kSeekEntryCount; ++i) {
2165     if (seek_head_.GetId(i) == kMkvCluster)
2166       cluster_index = i;
2167     if (seek_head_.GetId(i) == kMkvCues)
2168       cues_index = i;
2169   }
2170   seek_head_.SetSeekEntry(cues_index, kMkvCues,
2171                           seek_head_.GetPosition(cluster_index));
2172   seek_head_.SetSeekEntry(cluster_index, kMkvCluster,
2173                           cues_.Size() + seek_head_.GetPosition(cues_index));
2174 }
2175 
Init(IMkvWriter * ptr_writer)2176 bool Segment::Init(IMkvWriter* ptr_writer) {
2177   if (!ptr_writer) {
2178     return false;
2179   }
2180   writer_cluster_ = ptr_writer;
2181   writer_cues_ = ptr_writer;
2182   writer_header_ = ptr_writer;
2183   return segment_info_.Init();
2184 }
2185 
CopyAndMoveCuesBeforeClusters(mkvparser::IMkvReader * reader,IMkvWriter * writer)2186 bool Segment::CopyAndMoveCuesBeforeClusters(mkvparser::IMkvReader* reader,
2187                                             IMkvWriter* writer) {
2188   if (!writer->Seekable() || chunking_)
2189     return false;
2190   const int64 cluster_offset = cluster_list_[0]->size_position() -
2191                                GetUIntSize(kMkvCluster);
2192 
2193   // Copy the headers.
2194   if (!ChunkedCopy(reader, writer, 0, cluster_offset))
2195     return false;
2196 
2197   // Recompute cue positions and seek entries.
2198   MoveCuesBeforeClusters();
2199 
2200   // Write cues and seek entries.
2201   // TODO(vigneshv): As of now, it's safe to call seek_head_.Finalize() for the
2202   // second time with a different writer object. But the name Finalize() doesn't
2203   // indicate something we want to call more than once. So consider renaming it
2204   // to write() or some such.
2205   if (!cues_.Write(writer) || !seek_head_.Finalize(writer))
2206     return false;
2207 
2208   // Copy the Clusters.
2209   if (!ChunkedCopy(reader, writer, cluster_offset,
2210                    cluster_end_offset_ - cluster_offset))
2211     return false;
2212 
2213   // Update the Segment size in case the Cues size has changed.
2214   const int64 pos = writer->Position();
2215   const int64 segment_size = writer->Position() - payload_pos_;
2216   if (writer->Position(size_position_) ||
2217       WriteUIntSize(writer, segment_size, 8) ||
2218       writer->Position(pos))
2219     return false;
2220   return true;
2221 }
2222 
Finalize()2223 bool Segment::Finalize() {
2224   if (WriteFramesAll() < 0)
2225     return false;
2226 
2227   if (mode_ == kFile) {
2228     if (cluster_list_size_ > 0) {
2229       // Update last cluster's size
2230       Cluster* const old_cluster = cluster_list_[cluster_list_size_-1];
2231 
2232       if (!old_cluster || !old_cluster->Finalize())
2233         return false;
2234     }
2235 
2236     if (chunking_ && chunk_writer_cluster_) {
2237       chunk_writer_cluster_->Close();
2238       chunk_count_++;
2239     }
2240 
2241     const double duration =
2242         (static_cast<double>(last_timestamp_) + last_block_duration_) /
2243         segment_info_.timecode_scale();
2244     segment_info_.set_duration(duration);
2245     if (!segment_info_.Finalize(writer_header_))
2246       return false;
2247 
2248     if (output_cues_)
2249       if (!seek_head_.AddSeekEntry(kMkvCues, MaxOffset()))
2250         return false;
2251 
2252     if (chunking_) {
2253       if (!chunk_writer_cues_)
2254         return false;
2255 
2256       char* name = NULL;
2257       if (!UpdateChunkName("cues", &name))
2258         return false;
2259 
2260       const bool cues_open = chunk_writer_cues_->Open(name);
2261       delete [] name;
2262       if (!cues_open)
2263         return false;
2264     }
2265 
2266     cluster_end_offset_ = writer_cluster_->Position();
2267 
2268     // Write the seek headers and cues
2269     if (output_cues_)
2270       if (!cues_.Write(writer_cues_))
2271         return false;
2272 
2273     if (!seek_head_.Finalize(writer_header_))
2274       return false;
2275 
2276     if (writer_header_->Seekable()) {
2277       if (size_position_ == -1)
2278         return false;
2279 
2280       const int64 pos = writer_header_->Position();
2281       const int64 segment_size = MaxOffset();
2282 
2283       if (segment_size < 1)
2284         return false;
2285 
2286       if (writer_header_->Position(size_position_))
2287         return false;
2288 
2289       if (WriteUIntSize(writer_header_, segment_size, 8))
2290         return false;
2291 
2292       if (writer_header_->Position(pos))
2293         return false;
2294     }
2295 
2296     if (chunking_) {
2297       // Do not close any writers until the segment size has been written,
2298       // otherwise the size may be off.
2299       if (!chunk_writer_cues_ || !chunk_writer_header_)
2300         return false;
2301 
2302       chunk_writer_cues_->Close();
2303       chunk_writer_header_->Close();
2304     }
2305   }
2306 
2307   return true;
2308 }
2309 
AddTrack(int32 number)2310 Track* Segment::AddTrack(int32 number) {
2311   Track* const track = new (std::nothrow) Track(&seed_);  // NOLINT
2312 
2313   if (!track)
2314     return NULL;
2315 
2316   if (!tracks_.AddTrack(track, number)) {
2317     delete track;
2318     return NULL;
2319   }
2320 
2321   return track;
2322 }
2323 
AddChapter()2324 Chapter* Segment::AddChapter() {
2325   return chapters_.AddChapter(&seed_);
2326 }
2327 
AddVideoTrack(int32 width,int32 height,int32 number)2328 uint64 Segment::AddVideoTrack(int32 width, int32 height, int32 number) {
2329   VideoTrack* const track = new (std::nothrow) VideoTrack(&seed_);  // NOLINT
2330   if (!track)
2331     return 0;
2332 
2333   track->set_type(Tracks::kVideo);
2334   track->set_codec_id(Tracks::kVp8CodecId);
2335   track->set_width(width);
2336   track->set_height(height);
2337 
2338   tracks_.AddTrack(track, number);
2339   has_video_ = true;
2340 
2341   return track->number();
2342 }
2343 
AddCuePoint(uint64 timestamp,uint64 track)2344 bool Segment::AddCuePoint(uint64 timestamp, uint64 track) {
2345   if (cluster_list_size_  < 1)
2346     return false;
2347 
2348   const Cluster* const cluster = cluster_list_[cluster_list_size_-1];
2349   if (!cluster)
2350     return false;
2351 
2352   CuePoint* const cue = new (std::nothrow) CuePoint();  // NOLINT
2353   if (!cue)
2354     return false;
2355 
2356   cue->set_time(timestamp / segment_info_.timecode_scale());
2357   cue->set_block_number(cluster->blocks_added());
2358   cue->set_cluster_pos(cluster->position_for_cues());
2359   cue->set_track(track);
2360   if (!cues_.AddCue(cue))
2361     return false;
2362 
2363   new_cuepoint_ = false;
2364   return true;
2365 }
2366 
AddAudioTrack(int32 sample_rate,int32 channels,int32 number)2367 uint64 Segment::AddAudioTrack(int32 sample_rate,
2368                               int32 channels,
2369                               int32 number) {
2370   AudioTrack* const track = new (std::nothrow) AudioTrack(&seed_);  // NOLINT
2371   if (!track)
2372     return 0;
2373 
2374   track->set_type(Tracks::kAudio);
2375   track->set_codec_id(Tracks::kVorbisCodecId);
2376   track->set_sample_rate(sample_rate);
2377   track->set_channels(channels);
2378 
2379   tracks_.AddTrack(track, number);
2380 
2381   return track->number();
2382 }
2383 
AddFrame(const uint8 * frame,uint64 length,uint64 track_number,uint64 timestamp,bool is_key)2384 bool Segment::AddFrame(const uint8* frame,
2385                        uint64 length,
2386                        uint64 track_number,
2387                        uint64 timestamp,
2388                        bool is_key) {
2389   if (!frame)
2390     return false;
2391 
2392   if (!CheckHeaderInfo())
2393     return false;
2394 
2395   // Check for non-monotonically increasing timestamps.
2396   if (timestamp < last_timestamp_)
2397     return false;
2398 
2399   // If the segment has a video track hold onto audio frames to make sure the
2400   // audio that is associated with the start time of a video key-frame is
2401   // muxed into the same cluster.
2402   if (has_video_ && tracks_.TrackIsAudio(track_number) && !force_new_cluster_) {
2403     Frame* const new_frame = new (std::nothrow) Frame();
2404     if (new_frame == NULL || !new_frame->Init(frame, length))
2405       return false;
2406     new_frame->set_track_number(track_number);
2407     new_frame->set_timestamp(timestamp);
2408     new_frame->set_is_key(is_key);
2409 
2410     if (!QueueFrame(new_frame))
2411       return false;
2412 
2413     return true;
2414   }
2415 
2416   if (!DoNewClusterProcessing(track_number, timestamp, is_key))
2417     return false;
2418 
2419   if (cluster_list_size_ < 1)
2420     return false;
2421 
2422   Cluster* const cluster = cluster_list_[cluster_list_size_ - 1];
2423   if (!cluster)
2424     return false;
2425 
2426   const uint64 timecode_scale = segment_info_.timecode_scale();
2427   const uint64 abs_timecode = timestamp / timecode_scale;
2428 
2429   if (!cluster->AddFrame(frame,
2430                          length,
2431                          track_number,
2432                          abs_timecode,
2433                          is_key))
2434     return false;
2435 
2436   if (new_cuepoint_ && cues_track_ == track_number) {
2437     if (!AddCuePoint(timestamp, cues_track_))
2438       return false;
2439   }
2440 
2441   if (timestamp > last_timestamp_)
2442     last_timestamp_ = timestamp;
2443 
2444   return true;
2445 }
2446 
AddFrameWithAdditional(const uint8 * frame,uint64 length,const uint8 * additional,uint64 additional_length,uint64 add_id,uint64 track_number,uint64 timestamp,bool is_key)2447 bool Segment::AddFrameWithAdditional(const uint8* frame,
2448                                      uint64 length,
2449                                      const uint8* additional,
2450                                      uint64 additional_length,
2451                                      uint64 add_id,
2452                                      uint64 track_number,
2453                                      uint64 timestamp,
2454                                      bool is_key) {
2455   if (frame == NULL || additional == NULL)
2456     return false;
2457 
2458   if (!CheckHeaderInfo())
2459     return false;
2460 
2461   // Check for non-monotonically increasing timestamps.
2462   if (timestamp < last_timestamp_)
2463     return false;
2464 
2465   // If the segment has a video track hold onto audio frames to make sure the
2466   // audio that is associated with the start time of a video key-frame is
2467   // muxed into the same cluster.
2468   if (has_video_ && tracks_.TrackIsAudio(track_number) && !force_new_cluster_) {
2469     Frame* const new_frame = new (std::nothrow) Frame();
2470     if (new_frame == NULL || !new_frame->Init(frame, length))
2471       return false;
2472     new_frame->set_track_number(track_number);
2473     new_frame->set_timestamp(timestamp);
2474     new_frame->set_is_key(is_key);
2475 
2476     if (!QueueFrame(new_frame))
2477       return false;
2478 
2479     return true;
2480   }
2481 
2482   if (!DoNewClusterProcessing(track_number, timestamp, is_key))
2483     return false;
2484 
2485   if (cluster_list_size_ < 1)
2486     return false;
2487 
2488   Cluster* const cluster = cluster_list_[cluster_list_size_ - 1];
2489   if (cluster == NULL)
2490     return false;
2491 
2492   const uint64 timecode_scale = segment_info_.timecode_scale();
2493   const uint64 abs_timecode = timestamp / timecode_scale;
2494 
2495   if (!cluster->AddFrameWithAdditional(frame,
2496                                        length,
2497                                        additional,
2498                                        additional_length,
2499                                        add_id,
2500                                        track_number,
2501                                        abs_timecode,
2502                                        is_key))
2503     return false;
2504 
2505   if (new_cuepoint_ && cues_track_ == track_number) {
2506     if (!AddCuePoint(timestamp, cues_track_))
2507       return false;
2508   }
2509 
2510   if (timestamp > last_timestamp_)
2511     last_timestamp_ = timestamp;
2512 
2513   return true;
2514 }
2515 
AddFrameWithDiscardPadding(const uint8 * frame,uint64 length,int64 discard_padding,uint64 track_number,uint64 timestamp,bool is_key)2516 bool Segment::AddFrameWithDiscardPadding(const uint8* frame,
2517                                          uint64 length,
2518                                          int64 discard_padding,
2519                                          uint64 track_number,
2520                                          uint64 timestamp,
2521                                          bool is_key) {
2522   if (frame == NULL || discard_padding <= 0)
2523     return false;
2524 
2525   if (!CheckHeaderInfo())
2526     return false;
2527 
2528   // Check for non-monotonically increasing timestamps.
2529   if (timestamp < last_timestamp_)
2530     return false;
2531 
2532   // If the segment has a video track hold onto audio frames to make sure the
2533   // audio that is associated with the start time of a video key-frame is
2534   // muxed into the same cluster.
2535   if (has_video_ && tracks_.TrackIsAudio(track_number) && !force_new_cluster_) {
2536     Frame* const new_frame = new (std::nothrow) Frame();
2537     if (new_frame == NULL || !new_frame->Init(frame, length))
2538       return false;
2539     new_frame->set_track_number(track_number);
2540     new_frame->set_timestamp(timestamp);
2541     new_frame->set_is_key(is_key);
2542     new_frame->set_discard_padding(discard_padding);
2543 
2544     if (!QueueFrame(new_frame))
2545       return false;
2546 
2547     return true;
2548   }
2549 
2550   if (!DoNewClusterProcessing(track_number, timestamp, is_key))
2551     return false;
2552 
2553   if (cluster_list_size_ < 1)
2554     return false;
2555 
2556   Cluster* const cluster = cluster_list_[cluster_list_size_ - 1];
2557   if (!cluster)
2558     return false;
2559 
2560   const uint64 timecode_scale = segment_info_.timecode_scale();
2561   const uint64 abs_timecode = timestamp / timecode_scale;
2562 
2563   if (!cluster->AddFrameWithDiscardPadding(frame, length,
2564                                            discard_padding,
2565                                            track_number,
2566                                            abs_timecode,
2567                                            is_key)) {
2568     return false;
2569   }
2570 
2571   if (new_cuepoint_ && cues_track_ == track_number) {
2572     if (!AddCuePoint(timestamp, cues_track_))
2573       return false;
2574   }
2575 
2576   if (timestamp > last_timestamp_)
2577     last_timestamp_ = timestamp;
2578 
2579   return true;
2580 }
2581 
AddMetadata(const uint8 * frame,uint64 length,uint64 track_number,uint64 timestamp_ns,uint64 duration_ns)2582 bool Segment::AddMetadata(const uint8* frame,
2583                           uint64 length,
2584                           uint64 track_number,
2585                           uint64 timestamp_ns,
2586                           uint64 duration_ns) {
2587   if (!frame)
2588     return false;
2589 
2590   if (!CheckHeaderInfo())
2591     return false;
2592 
2593   // Check for non-monotonically increasing timestamps.
2594   if (timestamp_ns < last_timestamp_)
2595     return false;
2596 
2597   if (!DoNewClusterProcessing(track_number, timestamp_ns, true))
2598     return false;
2599 
2600   if (cluster_list_size_ < 1)
2601     return false;
2602 
2603   Cluster* const cluster = cluster_list_[cluster_list_size_-1];
2604 
2605   if (!cluster)
2606     return false;
2607 
2608   const uint64 timecode_scale = segment_info_.timecode_scale();
2609   const uint64 abs_timecode = timestamp_ns / timecode_scale;
2610   const uint64 duration_timecode = duration_ns / timecode_scale;
2611 
2612   if (!cluster->AddMetadata(frame,
2613                             length,
2614                             track_number,
2615                             abs_timecode,
2616                             duration_timecode))
2617     return false;
2618 
2619   if (timestamp_ns > last_timestamp_)
2620     last_timestamp_ = timestamp_ns;
2621 
2622   return true;
2623 }
2624 
AddGenericFrame(const Frame * frame)2625 bool Segment::AddGenericFrame(const Frame* frame) {
2626   last_block_duration_ = frame->duration();
2627   if (!tracks_.TrackIsAudio(frame->track_number()) &&
2628       !tracks_.TrackIsVideo(frame->track_number()) &&
2629       frame->duration() > 0) {
2630     return AddMetadata(frame->frame(),
2631                        frame->length(),
2632                        frame->track_number(),
2633                        frame->timestamp(),
2634                        frame->duration());
2635   } else if (frame->additional() && frame->additional_length() > 0) {
2636     return AddFrameWithAdditional(frame->frame(),
2637                                   frame->length(),
2638                                   frame->additional(),
2639                                   frame->additional_length(),
2640                                   frame->add_id(),
2641                                   frame->track_number(),
2642                                   frame->timestamp(),
2643                                   frame->is_key());
2644   } else if (frame->discard_padding() > 0) {
2645     return AddFrameWithDiscardPadding(frame->frame(), frame->length(),
2646                                       frame->discard_padding(),
2647                                       frame->track_number(),
2648                                       frame->timestamp(),
2649                                       frame->is_key());
2650   } else {
2651     return AddFrame(frame->frame(),
2652                     frame->length(),
2653                     frame->track_number(),
2654                     frame->timestamp(),
2655                     frame->is_key());
2656   }
2657 }
2658 
OutputCues(bool output_cues)2659 void Segment::OutputCues(bool output_cues) {
2660   output_cues_ = output_cues;
2661 }
2662 
SetChunking(bool chunking,const char * filename)2663 bool Segment::SetChunking(bool chunking, const char* filename) {
2664   if (chunk_count_ > 0)
2665     return false;
2666 
2667   if (chunking) {
2668     if (!filename)
2669       return false;
2670 
2671     // Check if we are being set to what is already set.
2672     if (chunking_ && !strcmp(filename, chunking_base_name_))
2673       return true;
2674 
2675     const size_t name_length = strlen(filename) + 1;
2676     char* const temp = new (std::nothrow) char[name_length];  // NOLINT
2677     if (!temp)
2678       return false;
2679 
2680 #ifdef _MSC_VER
2681     strcpy_s(temp, name_length, filename);
2682 #else
2683     strcpy(temp, filename);
2684 #endif
2685 
2686     delete [] chunking_base_name_;
2687     chunking_base_name_ = temp;
2688 
2689     if (!UpdateChunkName("chk", &chunk_name_))
2690       return false;
2691 
2692     if (!chunk_writer_cluster_) {
2693       chunk_writer_cluster_ = new (std::nothrow) MkvWriter();  // NOLINT
2694       if (!chunk_writer_cluster_)
2695         return false;
2696     }
2697 
2698     if (!chunk_writer_cues_) {
2699       chunk_writer_cues_ = new (std::nothrow) MkvWriter();  // NOLINT
2700       if (!chunk_writer_cues_)
2701         return false;
2702     }
2703 
2704     if (!chunk_writer_header_) {
2705       chunk_writer_header_ = new (std::nothrow) MkvWriter();  // NOLINT
2706       if (!chunk_writer_header_)
2707         return false;
2708     }
2709 
2710     if (!chunk_writer_cluster_->Open(chunk_name_))
2711       return false;
2712 
2713     const size_t header_length = strlen(filename) + strlen(".hdr") + 1;
2714     char* const header = new (std::nothrow) char[header_length];  // NOLINT
2715     if (!header)
2716       return false;
2717 
2718 #ifdef _MSC_VER
2719     strcpy_s(header, header_length - strlen(".hdr"), chunking_base_name_);
2720     strcat_s(header, header_length, ".hdr");
2721 #else
2722     strcpy(header, chunking_base_name_);
2723     strcat(header, ".hdr");
2724 #endif
2725     if (!chunk_writer_header_->Open(header)) {
2726       delete [] header;
2727       return false;
2728     }
2729 
2730     writer_cluster_ = chunk_writer_cluster_;
2731     writer_cues_ = chunk_writer_cues_;
2732     writer_header_ = chunk_writer_header_;
2733 
2734     delete [] header;
2735   }
2736 
2737   chunking_ = chunking;
2738 
2739   return true;
2740 }
2741 
CuesTrack(uint64 track_number)2742 bool Segment::CuesTrack(uint64 track_number) {
2743   const Track* const track = GetTrackByNumber(track_number);
2744   if (!track)
2745     return false;
2746 
2747   cues_track_ = track_number;
2748   return true;
2749 }
2750 
ForceNewClusterOnNextFrame()2751 void Segment::ForceNewClusterOnNextFrame() {
2752   force_new_cluster_ = true;
2753 }
2754 
GetTrackByNumber(uint64 track_number) const2755 Track* Segment::GetTrackByNumber(uint64 track_number) const {
2756   return tracks_.GetTrackByNumber(track_number);
2757 }
2758 
WriteSegmentHeader()2759 bool Segment::WriteSegmentHeader() {
2760   // TODO(fgalligan): Support more than one segment.
2761   if (!WriteEbmlHeader(writer_header_))
2762     return false;
2763 
2764   // Write "unknown" (-1) as segment size value. If mode is kFile, Segment
2765   // will write over duration when the file is finalized.
2766   if (WriteID(writer_header_, kMkvSegment))
2767     return false;
2768 
2769   // Save for later.
2770   size_position_ = writer_header_->Position();
2771 
2772   // Write "unknown" (EBML coded -1) as segment size value. We need to write 8
2773   // bytes because if we are going to overwrite the segment size later we do
2774   // not know how big our segment will be.
2775   if (SerializeInt(writer_header_, kEbmlUnknownValue, 8))
2776     return false;
2777 
2778   payload_pos_ =  writer_header_->Position();
2779 
2780   if (mode_ == kFile && writer_header_->Seekable()) {
2781     // Set the duration > 0.0 so SegmentInfo will write out the duration. When
2782     // the muxer is done writing we will set the correct duration and have
2783     // SegmentInfo upadte it.
2784     segment_info_.set_duration(1.0);
2785 
2786     if (!seek_head_.Write(writer_header_))
2787       return false;
2788   }
2789 
2790   if (!seek_head_.AddSeekEntry(kMkvInfo, MaxOffset()))
2791     return false;
2792   if (!segment_info_.Write(writer_header_))
2793     return false;
2794 
2795   if (!seek_head_.AddSeekEntry(kMkvTracks, MaxOffset()))
2796     return false;
2797   if (!tracks_.Write(writer_header_))
2798     return false;
2799 
2800   if (chapters_.Count() > 0) {
2801     if (!seek_head_.AddSeekEntry(kMkvChapters, MaxOffset()))
2802       return false;
2803     if (!chapters_.Write(writer_header_))
2804       return false;
2805   }
2806 
2807   if (chunking_ && (mode_ == kLive || !writer_header_->Seekable())) {
2808     if (!chunk_writer_header_)
2809       return false;
2810 
2811     chunk_writer_header_->Close();
2812   }
2813 
2814   header_written_ = true;
2815 
2816   return true;
2817 }
2818 
2819 // Here we are testing whether to create a new cluster, given a frame
2820 // having time frame_timestamp_ns.
2821 //
TestFrame(uint64 track_number,uint64 frame_timestamp_ns,bool is_key) const2822 int Segment::TestFrame(uint64 track_number,
2823                        uint64 frame_timestamp_ns,
2824                        bool is_key) const {
2825   if (force_new_cluster_)
2826     return 1;
2827 
2828   // If no clusters have been created yet, then create a new cluster
2829   // and write this frame immediately, in the new cluster.  This path
2830   // should only be followed once, the first time we attempt to write
2831   // a frame.
2832 
2833   if (cluster_list_size_ <= 0)
2834     return 1;
2835 
2836   // There exists at least one cluster. We must compare the frame to
2837   // the last cluster, in order to determine whether the frame is
2838   // written to the existing cluster, or that a new cluster should be
2839   // created.
2840 
2841   const uint64 timecode_scale = segment_info_.timecode_scale();
2842   const uint64 frame_timecode = frame_timestamp_ns / timecode_scale;
2843 
2844   const Cluster* const last_cluster = cluster_list_[cluster_list_size_ - 1];
2845   const uint64 last_cluster_timecode = last_cluster->timecode();
2846 
2847   // For completeness we test for the case when the frame's timecode
2848   // is less than the cluster's timecode.  Although in principle that
2849   // is allowed, this muxer doesn't actually write clusters like that,
2850   // so this indicates a bug somewhere in our algorithm.
2851 
2852   if (frame_timecode < last_cluster_timecode)  // should never happen
2853     return -1;  // error
2854 
2855   // If the frame has a timestamp significantly larger than the last
2856   // cluster (in Matroska, cluster-relative timestamps are serialized
2857   // using a 16-bit signed integer), then we cannot write this frame
2858   // to that cluster, and so we must create a new cluster.
2859 
2860   const int64 delta_timecode = frame_timecode - last_cluster_timecode;
2861 
2862   if (delta_timecode > kMaxBlockTimecode)
2863     return 2;
2864 
2865   // We decide to create a new cluster when we have a video keyframe.
2866   // This will flush queued (audio) frames, and write the keyframe
2867   // immediately, in the newly-created cluster.
2868 
2869   if (is_key && tracks_.TrackIsVideo(track_number))
2870     return 1;
2871 
2872   // Create a new cluster if we have accumulated too many frames
2873   // already, where "too many" is defined as "the total time of frames
2874   // in the cluster exceeds a threshold".
2875 
2876   const uint64 delta_ns = delta_timecode * timecode_scale;
2877 
2878   if (max_cluster_duration_ > 0 && delta_ns >= max_cluster_duration_)
2879     return 1;
2880 
2881   // This is similar to the case above, with the difference that a new
2882   // cluster is created when the size of the current cluster exceeds a
2883   // threshold.
2884 
2885   const uint64 cluster_size = last_cluster->payload_size();
2886 
2887   if (max_cluster_size_ > 0 && cluster_size >= max_cluster_size_)
2888     return 1;
2889 
2890   // There's no need to create a new cluster, so emit this frame now.
2891 
2892   return 0;
2893 }
2894 
MakeNewCluster(uint64 frame_timestamp_ns)2895 bool Segment::MakeNewCluster(uint64 frame_timestamp_ns) {
2896   const int32 new_size = cluster_list_size_ + 1;
2897 
2898   if (new_size > cluster_list_capacity_) {
2899     // Add more clusters.
2900     const int32 new_capacity =
2901         (cluster_list_capacity_ <= 0) ? 1 : cluster_list_capacity_ * 2;
2902     Cluster** const clusters =
2903         new (std::nothrow) Cluster*[new_capacity];  // NOLINT
2904     if (!clusters)
2905       return false;
2906 
2907     for (int32 i = 0; i < cluster_list_size_; ++i) {
2908       clusters[i] = cluster_list_[i];
2909     }
2910 
2911     delete [] cluster_list_;
2912 
2913     cluster_list_ = clusters;
2914     cluster_list_capacity_ = new_capacity;
2915   }
2916 
2917   if (!WriteFramesLessThan(frame_timestamp_ns))
2918     return false;
2919 
2920   if (mode_ == kFile) {
2921     if (cluster_list_size_ > 0) {
2922       // Update old cluster's size
2923       Cluster* const old_cluster = cluster_list_[cluster_list_size_ - 1];
2924 
2925       if (!old_cluster || !old_cluster->Finalize())
2926         return false;
2927     }
2928 
2929     if (output_cues_)
2930       new_cuepoint_ = true;
2931   }
2932 
2933   if (chunking_ && cluster_list_size_ > 0) {
2934     chunk_writer_cluster_->Close();
2935     chunk_count_++;
2936 
2937     if (!UpdateChunkName("chk", &chunk_name_))
2938       return false;
2939     if (!chunk_writer_cluster_->Open(chunk_name_))
2940       return false;
2941   }
2942 
2943   const uint64 timecode_scale = segment_info_.timecode_scale();
2944   const uint64 frame_timecode = frame_timestamp_ns / timecode_scale;
2945 
2946   uint64 cluster_timecode = frame_timecode;
2947 
2948   if (frames_size_ > 0) {
2949     const Frame* const f = frames_[0];  // earliest queued frame
2950     const uint64 ns = f->timestamp();
2951     const uint64 tc = ns / timecode_scale;
2952 
2953     if (tc < cluster_timecode)
2954       cluster_timecode = tc;
2955   }
2956 
2957   Cluster*& cluster = cluster_list_[cluster_list_size_];
2958   const int64 offset = MaxOffset();
2959   cluster = new (std::nothrow) Cluster(cluster_timecode, offset);  // NOLINT
2960   if (!cluster)
2961     return false;
2962 
2963   if (!cluster->Init(writer_cluster_))
2964     return false;
2965 
2966   cluster_list_size_ = new_size;
2967   return true;
2968 }
2969 
DoNewClusterProcessing(uint64 track_number,uint64 frame_timestamp_ns,bool is_key)2970 bool Segment::DoNewClusterProcessing(uint64 track_number,
2971                                      uint64 frame_timestamp_ns,
2972                                      bool is_key) {
2973   for (;;) {
2974     // Based on the characteristics of the current frame and current
2975     // cluster, decide whether to create a new cluster.
2976     const int result = TestFrame(track_number, frame_timestamp_ns, is_key);
2977     if (result < 0)  // error
2978       return false;
2979 
2980   // Always set force_new_cluster_ to false after TestFrame.
2981   force_new_cluster_ = false;
2982 
2983   // A non-zero result means create a new cluster.
2984   if (result > 0 && !MakeNewCluster(frame_timestamp_ns))
2985     return false;
2986 
2987     // Write queued (audio) frames.
2988     const int frame_count = WriteFramesAll();
2989     if (frame_count < 0)  // error
2990       return false;
2991 
2992     // Write the current frame to the current cluster (if TestFrame
2993     // returns 0) or to a newly created cluster (TestFrame returns 1).
2994     if (result <= 1)
2995       return true;
2996 
2997     // TestFrame returned 2, which means there was a large time
2998     // difference between the cluster and the frame itself.  Do the
2999     // test again, comparing the frame to the new cluster.
3000   }
3001 }
3002 
CheckHeaderInfo()3003 bool Segment::CheckHeaderInfo() {
3004   if (!header_written_) {
3005     if (!WriteSegmentHeader())
3006       return false;
3007 
3008     if (!seek_head_.AddSeekEntry(kMkvCluster, MaxOffset()))
3009       return false;
3010 
3011     if (output_cues_ && cues_track_ == 0) {
3012       // Check for a video track
3013       for (uint32 i = 0; i < tracks_.track_entries_size(); ++i) {
3014         const Track* const track = tracks_.GetTrackByIndex(i);
3015         if (!track)
3016           return false;
3017 
3018         if (tracks_.TrackIsVideo(track->number())) {
3019           cues_track_ = track->number();
3020           break;
3021         }
3022       }
3023 
3024       // Set first track found
3025       if (cues_track_ == 0) {
3026         const Track* const track = tracks_.GetTrackByIndex(0);
3027         if (!track)
3028           return false;
3029 
3030         cues_track_ = track->number();
3031       }
3032     }
3033   }
3034   return true;
3035 }
3036 
UpdateChunkName(const char * ext,char ** name) const3037 bool Segment::UpdateChunkName(const char* ext, char** name) const {
3038   if (!name || !ext)
3039     return false;
3040 
3041   char ext_chk[64];
3042 #ifdef _MSC_VER
3043   sprintf_s(ext_chk, sizeof(ext_chk), "_%06d.%s", chunk_count_, ext);
3044 #else
3045   snprintf(ext_chk, sizeof(ext_chk), "_%06d.%s", chunk_count_, ext);
3046 #endif
3047 
3048   const size_t length = strlen(chunking_base_name_) + strlen(ext_chk) + 1;
3049   char* const str = new (std::nothrow) char[length];  // NOLINT
3050   if (!str)
3051     return false;
3052 
3053 #ifdef _MSC_VER
3054   strcpy_s(str, length-strlen(ext_chk), chunking_base_name_);
3055   strcat_s(str, length, ext_chk);
3056 #else
3057   strcpy(str, chunking_base_name_);
3058   strcat(str, ext_chk);
3059 #endif
3060 
3061   delete [] *name;
3062   *name = str;
3063 
3064   return true;
3065 }
3066 
MaxOffset()3067 int64 Segment::MaxOffset() {
3068   if (!writer_header_)
3069     return -1;
3070 
3071   int64 offset = writer_header_->Position() - payload_pos_;
3072 
3073   if (chunking_) {
3074     for (int32 i = 0; i < cluster_list_size_; ++i) {
3075       Cluster* const cluster = cluster_list_[i];
3076       offset += cluster->Size();
3077     }
3078 
3079     if (writer_cues_)
3080       offset += writer_cues_->Position();
3081   }
3082 
3083   return offset;
3084 }
3085 
QueueFrame(Frame * frame)3086 bool Segment::QueueFrame(Frame* frame) {
3087   const int32 new_size = frames_size_ + 1;
3088 
3089   if (new_size > frames_capacity_) {
3090     // Add more frames.
3091     const int32 new_capacity = (!frames_capacity_) ? 2 : frames_capacity_ * 2;
3092 
3093     if (new_capacity < 1)
3094       return false;
3095 
3096     Frame** const frames = new (std::nothrow) Frame*[new_capacity];  // NOLINT
3097     if (!frames)
3098       return false;
3099 
3100     for (int32 i = 0; i < frames_size_; ++i) {
3101       frames[i] = frames_[i];
3102     }
3103 
3104     delete [] frames_;
3105     frames_ = frames;
3106     frames_capacity_ = new_capacity;
3107   }
3108 
3109   frames_[frames_size_++] = frame;
3110 
3111   return true;
3112 }
3113 
WriteFramesAll()3114 int Segment::WriteFramesAll() {
3115   if (frames_ == NULL)
3116     return 0;
3117 
3118   if (cluster_list_size_ < 1)
3119     return -1;
3120 
3121   Cluster* const cluster = cluster_list_[cluster_list_size_-1];
3122 
3123   if (!cluster)
3124     return -1;
3125 
3126   const uint64 timecode_scale = segment_info_.timecode_scale();
3127 
3128   for (int32 i = 0; i < frames_size_; ++i) {
3129     Frame*& frame = frames_[i];
3130     const uint64 frame_timestamp = frame->timestamp();  // ns
3131     const uint64 frame_timecode = frame_timestamp / timecode_scale;
3132 
3133     if (frame->discard_padding() > 0) {
3134       if (!cluster->AddFrameWithDiscardPadding(frame->frame(),
3135                                                frame->length(),
3136                                                frame->discard_padding(),
3137                                                frame->track_number(),
3138                                                frame_timecode,
3139                                                frame->is_key())) {
3140         return -1;
3141       }
3142     } else {
3143       if (!cluster->AddFrame(frame->frame(),
3144                              frame->length(),
3145                              frame->track_number(),
3146                              frame_timecode,
3147                              frame->is_key())) {
3148         return -1;
3149       }
3150     }
3151 
3152     if (new_cuepoint_ && cues_track_ == frame->track_number()) {
3153       if (!AddCuePoint(frame_timestamp, cues_track_))
3154         return -1;
3155     }
3156 
3157     if (frame_timestamp > last_timestamp_)
3158       last_timestamp_ = frame_timestamp;
3159 
3160     delete frame;
3161     frame = NULL;
3162   }
3163 
3164   const int result = frames_size_;
3165   frames_size_ = 0;
3166 
3167   return result;
3168 }
3169 
WriteFramesLessThan(uint64 timestamp)3170 bool Segment::WriteFramesLessThan(uint64 timestamp) {
3171   // Check |cluster_list_size_| to see if this is the first cluster. If it is
3172   // the first cluster the audio frames that are less than the first video
3173   // timesatmp will be written in a later step.
3174   if (frames_size_ > 0 && cluster_list_size_ > 0) {
3175     if (!frames_)
3176       return false;
3177 
3178     Cluster* const cluster = cluster_list_[cluster_list_size_-1];
3179     if (!cluster)
3180       return false;
3181 
3182     const uint64 timecode_scale = segment_info_.timecode_scale();
3183     int32 shift_left = 0;
3184 
3185     // TODO(fgalligan): Change this to use the durations of frames instead of
3186     // the next frame's start time if the duration is accurate.
3187     for (int32 i = 1; i < frames_size_; ++i) {
3188       const Frame* const frame_curr = frames_[i];
3189 
3190       if (frame_curr->timestamp() > timestamp)
3191         break;
3192 
3193       const Frame* const frame_prev = frames_[i-1];
3194       const uint64 frame_timestamp = frame_prev->timestamp();
3195       const uint64 frame_timecode = frame_timestamp / timecode_scale;
3196       const int64 discard_padding = frame_prev->discard_padding();
3197 
3198       if (discard_padding > 0) {
3199         if (!cluster->AddFrameWithDiscardPadding(frame_prev->frame(),
3200                                                  frame_prev->length(),
3201                                                  discard_padding,
3202                                                  frame_prev->track_number(),
3203                                                  frame_timecode,
3204                                                  frame_prev->is_key())) {
3205           return false;
3206         }
3207       } else {
3208         if (!cluster->AddFrame(frame_prev->frame(),
3209                                frame_prev->length(),
3210                                frame_prev->track_number(),
3211                                frame_timecode,
3212                                frame_prev->is_key())) {
3213           return false;
3214         }
3215       }
3216 
3217       if (new_cuepoint_ && cues_track_ == frame_prev->track_number()) {
3218         if (!AddCuePoint(frame_timestamp, cues_track_))
3219           return false;
3220       }
3221 
3222       ++shift_left;
3223       if (frame_timestamp > last_timestamp_)
3224         last_timestamp_ = frame_timestamp;
3225 
3226       delete frame_prev;
3227     }
3228 
3229     if (shift_left > 0) {
3230       if (shift_left >= frames_size_)
3231         return false;
3232 
3233       const int32 new_frames_size = frames_size_ - shift_left;
3234       for (int32 i = 0; i < new_frames_size; ++i) {
3235         frames_[i] = frames_[i+shift_left];
3236       }
3237 
3238       frames_size_ = new_frames_size;
3239     }
3240   }
3241 
3242   return true;
3243 }
3244 
3245 }  // namespace mkvmuxer
3246