mp4/
track.rs

1use bytes::BytesMut;
2use std::cmp;
3use std::convert::TryFrom;
4use std::io::{Read, Seek, SeekFrom, Write};
5use std::time::Duration;
6
7use crate::mp4box::traf::TrafBox;
8use crate::mp4box::trak::TrakBox;
9use crate::mp4box::trun::TrunBox;
10use crate::mp4box::{
11    avc1::Avc1Box, co64::Co64Box, ctts::CttsBox, ctts::CttsEntry, hev1::Hev1Box, mp4a::Mp4aBox,
12    smhd::SmhdBox, stco::StcoBox, stsc::StscEntry, stss::StssBox, stts::SttsEntry, tx3g::Tx3gBox,
13    vmhd::VmhdBox, vp09::Vp09Box,
14};
15use crate::*;
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct TrackConfig {
19    pub track_type: TrackType,
20    pub timescale: u32,
21    pub language: String,
22    pub media_conf: MediaConfig,
23}
24
25impl From<MediaConfig> for TrackConfig {
26    fn from(media_conf: MediaConfig) -> Self {
27        match media_conf {
28            MediaConfig::AvcConfig(avc_conf) => Self::from(avc_conf),
29            MediaConfig::HevcConfig(hevc_conf) => Self::from(hevc_conf),
30            MediaConfig::AacConfig(aac_conf) => Self::from(aac_conf),
31            MediaConfig::TtxtConfig(ttxt_conf) => Self::from(ttxt_conf),
32            MediaConfig::Vp9Config(vp9_config) => Self::from(vp9_config),
33        }
34    }
35}
36
37impl From<AvcConfig> for TrackConfig {
38    fn from(avc_conf: AvcConfig) -> Self {
39        Self {
40            track_type: TrackType::Video,
41            timescale: 1000,               // XXX
42            language: String::from("und"), // XXX
43            media_conf: MediaConfig::AvcConfig(avc_conf),
44        }
45    }
46}
47
48impl From<HevcConfig> for TrackConfig {
49    fn from(hevc_conf: HevcConfig) -> Self {
50        Self {
51            track_type: TrackType::Video,
52            timescale: 1000,               // XXX
53            language: String::from("und"), // XXX
54            media_conf: MediaConfig::HevcConfig(hevc_conf),
55        }
56    }
57}
58
59impl From<AacConfig> for TrackConfig {
60    fn from(aac_conf: AacConfig) -> Self {
61        Self {
62            track_type: TrackType::Audio,
63            timescale: 1000,               // XXX
64            language: String::from("und"), // XXX
65            media_conf: MediaConfig::AacConfig(aac_conf),
66        }
67    }
68}
69
70impl From<TtxtConfig> for TrackConfig {
71    fn from(txtt_conf: TtxtConfig) -> Self {
72        Self {
73            track_type: TrackType::Subtitle,
74            timescale: 1000,               // XXX
75            language: String::from("und"), // XXX
76            media_conf: MediaConfig::TtxtConfig(txtt_conf),
77        }
78    }
79}
80
81impl From<Vp9Config> for TrackConfig {
82    fn from(vp9_conf: Vp9Config) -> Self {
83        Self {
84            track_type: TrackType::Video,
85            timescale: 1000,               // XXX
86            language: String::from("und"), // XXX
87            media_conf: MediaConfig::Vp9Config(vp9_conf),
88        }
89    }
90}
91
92fn sample_times(trak: &TrakBox) -> Result<Vec<(u64, u32)>> {
93    let stts = &trak.mdia.minf.stbl.stts;
94
95    let mut sample_count: u32 = 1;
96    let mut elapsed = 0;
97
98    let mut result = Vec::new();
99
100    let mut sample_id = 1;
101
102    for entry in stts.entries.iter() {
103        let new_sample_count =
104            sample_count
105                .checked_add(entry.sample_count)
106                .ok_or(Error::InvalidData(
107                    "attempt to sum stts entries sample_count with overflow",
108                ))?;
109
110        while sample_id < new_sample_count {
111            let start_time =
112                (sample_id - sample_count) as u64 * entry.sample_delta as u64 + elapsed;
113            sample_id += 1;
114            result.push((start_time, entry.sample_delta));
115        }
116
117        sample_count = new_sample_count;
118        elapsed += entry.sample_count as u64 * entry.sample_delta as u64;
119    }
120
121    Ok(result)
122}
123
124#[derive(Debug)]
125pub struct Mp4Track {
126    pub trak: TrakBox,
127    pub trafs: Vec<TrafBox>,
128    pub moof_offsets: Vec<u64>,
129
130    // Fragmented Tracks Defaults.
131    pub default_sample_duration: u32,
132
133    sample_time_cache: Vec<(u64, u32)>,
134}
135
136impl Mp4Track {
137    pub(crate) fn from(trak: &TrakBox) -> Self {
138        let trak = trak.clone();
139        let sample_time_cache = sample_times(&trak).unwrap();
140        Self {
141            trak,
142            trafs: Vec::new(),
143            moof_offsets: Vec::new(),
144            default_sample_duration: 0,
145            sample_time_cache,
146        }
147    }
148
149    pub fn track_id(&self) -> u32 {
150        self.trak.tkhd.track_id
151    }
152
153    pub fn track_type(&self) -> Result<TrackType> {
154        TrackType::try_from(&self.trak.mdia.hdlr.handler_type)
155    }
156
157    pub fn media_type(&self) -> Result<MediaType> {
158        if self.trak.mdia.minf.stbl.stsd.avc1.is_some() {
159            Ok(MediaType::H264)
160        } else if self.trak.mdia.minf.stbl.stsd.hev1.is_some() {
161            Ok(MediaType::H265)
162        } else if self.trak.mdia.minf.stbl.stsd.vp09.is_some() {
163            Ok(MediaType::VP9)
164        } else if self.trak.mdia.minf.stbl.stsd.mp4a.is_some() {
165            Ok(MediaType::AAC)
166        } else if self.trak.mdia.minf.stbl.stsd.tx3g.is_some() {
167            Ok(MediaType::TTXT)
168        } else {
169            Err(Error::InvalidData("unsupported media type"))
170        }
171    }
172
173    pub fn box_type(&self) -> Result<FourCC> {
174        if self.trak.mdia.minf.stbl.stsd.avc1.is_some() {
175            Ok(FourCC::from(BoxType::Avc1Box))
176        } else if self.trak.mdia.minf.stbl.stsd.hev1.is_some() {
177            Ok(FourCC::from(BoxType::Hev1Box))
178        } else if self.trak.mdia.minf.stbl.stsd.vp09.is_some() {
179            Ok(FourCC::from(BoxType::Vp09Box))
180        } else if self.trak.mdia.minf.stbl.stsd.mp4a.is_some() {
181            Ok(FourCC::from(BoxType::Mp4aBox))
182        } else if self.trak.mdia.minf.stbl.stsd.tx3g.is_some() {
183            Ok(FourCC::from(BoxType::Tx3gBox))
184        } else {
185            Err(Error::InvalidData("unsupported sample entry box"))
186        }
187    }
188
189    pub fn width(&self) -> u16 {
190        if let Some(ref avc1) = self.trak.mdia.minf.stbl.stsd.avc1 {
191            avc1.width
192        } else {
193            self.trak.tkhd.width.value()
194        }
195    }
196
197    pub fn height(&self) -> u16 {
198        if let Some(ref avc1) = self.trak.mdia.minf.stbl.stsd.avc1 {
199            avc1.height
200        } else {
201            self.trak.tkhd.height.value()
202        }
203    }
204
205    pub fn frame_rate(&self) -> f64 {
206        let dur = self.duration();
207        if dur.is_zero() {
208            0.0
209        } else {
210            self.sample_count() as f64 / dur.as_secs_f64()
211        }
212    }
213
214    pub fn sample_freq_index(&self) -> Result<SampleFreqIndex> {
215        if let Some(ref mp4a) = self.trak.mdia.minf.stbl.stsd.mp4a {
216            if let Some(ref esds) = mp4a.esds {
217                SampleFreqIndex::try_from(esds.es_desc.dec_config.dec_specific.freq_index)
218            } else {
219                Err(Error::BoxInStblNotFound(self.track_id(), BoxType::EsdsBox))
220            }
221        } else {
222            Err(Error::BoxInStblNotFound(self.track_id(), BoxType::Mp4aBox))
223        }
224    }
225
226    pub fn channel_config(&self) -> Result<ChannelConfig> {
227        if let Some(ref mp4a) = self.trak.mdia.minf.stbl.stsd.mp4a {
228            if let Some(ref esds) = mp4a.esds {
229                ChannelConfig::try_from(esds.es_desc.dec_config.dec_specific.chan_conf)
230            } else {
231                Err(Error::BoxInStblNotFound(self.track_id(), BoxType::EsdsBox))
232            }
233        } else {
234            Err(Error::BoxInStblNotFound(self.track_id(), BoxType::Mp4aBox))
235        }
236    }
237
238    pub fn language(&self) -> &str {
239        &self.trak.mdia.mdhd.language
240    }
241
242    pub fn timescale(&self) -> u32 {
243        self.trak.mdia.mdhd.timescale
244    }
245
246    pub fn duration(&self) -> Duration {
247        Duration::from_micros(
248            self.trak.mdia.mdhd.duration * 1_000_000 / self.trak.mdia.mdhd.timescale as u64,
249        )
250    }
251
252    pub fn bitrate(&self) -> u32 {
253        if let Some(ref mp4a) = self.trak.mdia.minf.stbl.stsd.mp4a {
254            if let Some(ref esds) = mp4a.esds {
255                esds.es_desc.dec_config.avg_bitrate
256            } else {
257                0
258            }
259            // mp4a.esds.es_desc.dec_config.avg_bitrate
260        } else {
261            let dur = self.duration();
262            if dur.is_zero() {
263                0
264            } else {
265                let bitrate = self.total_sample_size() as f64 * 8.0 / dur.as_secs_f64();
266                bitrate as u32
267            }
268        }
269    }
270
271    pub fn sample_count(&self) -> u32 {
272        if !self.trafs.is_empty() {
273            let mut sample_count = 0u32;
274            for traf in self.trafs.iter() {
275                if let Some(ref trun) = traf.trun {
276                    sample_count = sample_count
277                        .checked_add(trun.sample_count)
278                        .expect("attempt to sum trun sample_count with overflow");
279                }
280            }
281            sample_count
282        } else {
283            self.trak.mdia.minf.stbl.stsz.sample_count
284        }
285    }
286
287    pub fn video_profile(&self) -> Result<AvcProfile> {
288        if let Some(ref avc1) = self.trak.mdia.minf.stbl.stsd.avc1 {
289            AvcProfile::try_from((
290                avc1.avcc.avc_profile_indication,
291                avc1.avcc.profile_compatibility,
292            ))
293        } else {
294            Err(Error::BoxInStblNotFound(self.track_id(), BoxType::Avc1Box))
295        }
296    }
297
298    pub fn sequence_parameter_set(&self) -> Result<&[u8]> {
299        if let Some(ref avc1) = self.trak.mdia.minf.stbl.stsd.avc1 {
300            match avc1.avcc.sequence_parameter_sets.get(0) {
301                Some(nal) => Ok(nal.bytes.as_ref()),
302                None => Err(Error::EntryInStblNotFound(
303                    self.track_id(),
304                    BoxType::AvcCBox,
305                    0,
306                )),
307            }
308        } else {
309            Err(Error::BoxInStblNotFound(self.track_id(), BoxType::Avc1Box))
310        }
311    }
312
313    pub fn picture_parameter_set(&self) -> Result<&[u8]> {
314        if let Some(ref avc1) = self.trak.mdia.minf.stbl.stsd.avc1 {
315            match avc1.avcc.picture_parameter_sets.get(0) {
316                Some(nal) => Ok(nal.bytes.as_ref()),
317                None => Err(Error::EntryInStblNotFound(
318                    self.track_id(),
319                    BoxType::AvcCBox,
320                    0,
321                )),
322            }
323        } else {
324            Err(Error::BoxInStblNotFound(self.track_id(), BoxType::Avc1Box))
325        }
326    }
327
328    pub fn audio_profile(&self) -> Result<AudioObjectType> {
329        if let Some(ref mp4a) = self.trak.mdia.minf.stbl.stsd.mp4a {
330            if let Some(ref esds) = mp4a.esds {
331                AudioObjectType::try_from(esds.es_desc.dec_config.dec_specific.profile)
332            } else {
333                Err(Error::BoxInStblNotFound(self.track_id(), BoxType::EsdsBox))
334            }
335        } else {
336            Err(Error::BoxInStblNotFound(self.track_id(), BoxType::Mp4aBox))
337        }
338    }
339
340    fn stsc_index(&self, sample_id: u32) -> Result<usize> {
341        if self.trak.mdia.minf.stbl.stsc.entries.is_empty() {
342            return Err(Error::InvalidData("no stsc entries"));
343        }
344        for (i, entry) in self.trak.mdia.minf.stbl.stsc.entries.iter().enumerate() {
345            if sample_id < entry.first_sample {
346                return if i == 0 {
347                    Err(Error::InvalidData("sample not found"))
348                } else {
349                    Ok(i - 1)
350                };
351            }
352        }
353        Ok(self.trak.mdia.minf.stbl.stsc.entries.len() - 1)
354    }
355
356    fn chunk_offset(&self, chunk_id: u32) -> Result<u64> {
357        if self.trak.mdia.minf.stbl.stco.is_none() && self.trak.mdia.minf.stbl.co64.is_none() {
358            return Err(Error::InvalidData("must have either stco or co64 boxes"));
359        }
360        if let Some(ref stco) = self.trak.mdia.minf.stbl.stco {
361            if let Some(offset) = stco.entries.get(chunk_id as usize - 1) {
362                return Ok(*offset as u64);
363            } else {
364                return Err(Error::EntryInStblNotFound(
365                    self.track_id(),
366                    BoxType::StcoBox,
367                    chunk_id,
368                ));
369            }
370        } else if let Some(ref co64) = self.trak.mdia.minf.stbl.co64 {
371            if let Some(offset) = co64.entries.get(chunk_id as usize - 1) {
372                return Ok(*offset);
373            } else {
374                return Err(Error::EntryInStblNotFound(
375                    self.track_id(),
376                    BoxType::Co64Box,
377                    chunk_id,
378                ));
379            }
380        }
381        Err(Error::Box2NotFound(BoxType::StcoBox, BoxType::Co64Box))
382    }
383
384    fn ctts_index(&self, sample_id: u32) -> Result<(usize, u32)> {
385        let ctts = self.trak.mdia.minf.stbl.ctts.as_ref().unwrap();
386        let mut sample_count: u32 = 1;
387        for (i, entry) in ctts.entries.iter().enumerate() {
388            let next_sample_count =
389                sample_count
390                    .checked_add(entry.sample_count)
391                    .ok_or(Error::InvalidData(
392                        "attempt to sum ctts entries sample_count with overflow",
393                    ))?;
394            if sample_id < next_sample_count {
395                return Ok((i, sample_count));
396            }
397            sample_count = next_sample_count;
398        }
399
400        Err(Error::EntryInStblNotFound(
401            self.track_id(),
402            BoxType::CttsBox,
403            sample_id,
404        ))
405    }
406
407    /// return `(traf_idx, sample_idx_in_trun)`
408    fn find_traf_idx_and_sample_idx(&self, sample_id: u32) -> Option<(usize, usize)> {
409        let global_idx = sample_id - 1;
410        let mut offset = 0;
411        for traf_idx in 0..self.trafs.len() {
412            if let Some(trun) = &self.trafs[traf_idx].trun {
413                let sample_count = trun.sample_count;
414                if sample_count > (global_idx - offset) {
415                    return Some((traf_idx, (global_idx - offset) as _));
416                }
417                offset = offset
418                    .checked_add(sample_count)
419                    .expect("attempt to sum trun sample_count with overflow");
420            }
421        }
422        None
423    }
424
425    fn sample_size(&self, sample_id: u32) -> Result<u32> {
426        if !self.trafs.is_empty() {
427            if let Some((traf_idx, sample_idx)) = self.find_traf_idx_and_sample_idx(sample_id) {
428                if let Some(size) = self.trafs[traf_idx]
429                    .trun
430                    .as_ref()
431                    .unwrap()
432                    .sample_sizes
433                    .get(sample_idx)
434                {
435                    Ok(*size)
436                } else {
437                    Err(Error::EntryInTrunNotFound(
438                        self.track_id(),
439                        BoxType::TrunBox,
440                        sample_id,
441                    ))
442                }
443            } else {
444                Err(Error::BoxInTrafNotFound(self.track_id(), BoxType::TrafBox))
445            }
446        } else {
447            let stsz = &self.trak.mdia.minf.stbl.stsz;
448            if stsz.sample_size > 0 {
449                return Ok(stsz.sample_size);
450            }
451            if let Some(size) = stsz.sample_sizes.get(sample_id as usize - 1) {
452                Ok(*size)
453            } else {
454                Err(Error::EntryInStblNotFound(
455                    self.track_id(),
456                    BoxType::StszBox,
457                    sample_id,
458                ))
459            }
460        }
461    }
462
463    fn total_sample_size(&self) -> u64 {
464        let stsz = &self.trak.mdia.minf.stbl.stsz;
465        if stsz.sample_size > 0 {
466            stsz.sample_size as u64 * self.sample_count() as u64
467        } else {
468            let mut total_size = 0;
469            for size in stsz.sample_sizes.iter() {
470                total_size += *size as u64;
471            }
472            total_size
473        }
474    }
475
476    pub fn sample_offset(&self, sample_id: u32) -> Result<u64> {
477        if !self.trafs.is_empty() {
478            if let Some((traf_idx, sample_idx)) = self.find_traf_idx_and_sample_idx(sample_id) {
479                let mut sample_offset = self.trafs[traf_idx]
480                    .tfhd
481                    .base_data_offset
482                    .unwrap_or(self.moof_offsets[traf_idx]);
483
484                if let Some(data_offset) = self.trafs[traf_idx]
485                    .trun
486                    .as_ref()
487                    .and_then(|trun| trun.data_offset)
488                {
489                    sample_offset = sample_offset.checked_add_signed(data_offset as i64).ok_or(
490                        Error::InvalidData("attempt to calculate trun sample offset with overflow"),
491                    )?;
492                }
493
494                let first_sample_in_trun = sample_id - sample_idx as u32;
495                for i in first_sample_in_trun..sample_id {
496                    sample_offset = sample_offset
497                        .checked_add(self.sample_size(i)? as u64)
498                        .ok_or(Error::InvalidData(
499                            "attempt to calculate trun entry sample offset with overflow",
500                        ))?;
501                }
502
503                Ok(sample_offset)
504            } else {
505                Err(Error::BoxInTrafNotFound(self.track_id(), BoxType::TrafBox))
506            }
507        } else {
508            let stsc_index = self.stsc_index(sample_id)?;
509
510            let stsc = &self.trak.mdia.minf.stbl.stsc;
511            let stsc_entry = stsc.entries.get(stsc_index).unwrap();
512
513            let first_chunk = stsc_entry.first_chunk;
514            let first_sample = stsc_entry.first_sample;
515            let samples_per_chunk = stsc_entry.samples_per_chunk;
516
517            let chunk_id = sample_id
518                .checked_sub(first_sample)
519                .map(|n| n / samples_per_chunk)
520                .and_then(|n| n.checked_add(first_chunk))
521                .ok_or(Error::InvalidData(
522                    "attempt to calculate stsc chunk_id with overflow",
523                ))?;
524
525            let chunk_offset = self.chunk_offset(chunk_id)?;
526
527            let first_sample_in_chunk = sample_id - (sample_id - first_sample) % samples_per_chunk;
528
529            let mut sample_offset: u64 = 0;
530            for i in first_sample_in_chunk..sample_id {
531                sample_offset += self.sample_size(i)? as u64;
532            }
533
534            Ok(chunk_offset + sample_offset)
535        }
536    }
537
538    pub(crate) fn sample_time(&self, sample_id: u32) -> Result<(u64, u32)> {
539        if !self.trafs.is_empty() {
540            let mut base_start_time = 0;
541            let mut default_sample_duration = self.default_sample_duration;
542            if let Some((traf_idx, sample_idx)) = self.find_traf_idx_and_sample_idx(sample_id) {
543                let traf = &self.trafs[traf_idx];
544                if let Some(tfdt) = &traf.tfdt {
545                    base_start_time = tfdt.base_media_decode_time;
546                }
547                if let Some(duration) = traf.tfhd.default_sample_duration {
548                    default_sample_duration = duration;
549                }
550                if let Some(trun) = &traf.trun {
551                    if TrunBox::FLAG_SAMPLE_DURATION & trun.flags != 0 {
552                        let mut start_offset = 0u64;
553                        for duration in &trun.sample_durations[..sample_idx] {
554                            start_offset = start_offset.checked_add(*duration as u64).ok_or(
555                                Error::InvalidData("attempt to sum sample durations with overflow"),
556                            )?;
557                        }
558                        let duration = trun.sample_durations[sample_idx];
559                        return Ok((base_start_time + start_offset, duration));
560                    }
561                }
562            }
563            let start_offset = ((sample_id - 1) * default_sample_duration) as u64;
564            Ok((base_start_time + start_offset, default_sample_duration))
565        } else {
566            Ok(self.sample_time_cache[sample_id as usize - 1])
567        }
568    }
569
570    fn sample_rendering_offset(&self, sample_id: u32) -> i32 {
571        if !self.trafs.is_empty() {
572            if let Some((traf_idx, sample_idx)) = self.find_traf_idx_and_sample_idx(sample_id) {
573                if let Some(cts) = self.trafs[traf_idx]
574                    .trun
575                    .as_ref()
576                    .and_then(|trun| trun.sample_cts.get(sample_idx))
577                {
578                    return *cts as i32;
579                }
580            }
581        } else if let Some(ref ctts) = self.trak.mdia.minf.stbl.ctts {
582            if let Ok((ctts_index, _)) = self.ctts_index(sample_id) {
583                let ctts_entry = ctts.entries.get(ctts_index).unwrap();
584                return ctts_entry.sample_offset;
585            }
586        }
587        0
588    }
589
590    fn is_sync_sample(&self, sample_id: u32) -> bool {
591        if !self.trafs.is_empty() {
592            let sample_sizes_count = self.sample_count() / self.trafs.len() as u32;
593            return sample_id == 1 || sample_id % sample_sizes_count == 0;
594        }
595
596        if let Some(ref stss) = self.trak.mdia.minf.stbl.stss {
597            stss.entries.binary_search(&sample_id).is_ok()
598        } else {
599            true
600        }
601    }
602
603    pub(crate) fn read_sample<R: Read + Seek>(
604        &self,
605        reader: &mut R,
606        sample_id: u32,
607    ) -> Result<Option<Mp4Sample>> {
608        let sample_offset = match self.sample_offset(sample_id) {
609            Ok(offset) => offset,
610            Err(Error::EntryInStblNotFound(_, _, _)) => return Ok(None),
611            Err(err) => return Err(err),
612        };
613        let sample_size = match self.sample_size(sample_id) {
614            Ok(size) => size,
615            Err(Error::EntryInStblNotFound(_, _, _)) => return Ok(None),
616            Err(err) => return Err(err),
617        };
618
619        let mut buffer = vec![0x0u8; sample_size as usize];
620        reader.seek(SeekFrom::Start(sample_offset))?;
621        reader.read_exact(&mut buffer)?;
622
623        let (start_time, duration) = self.sample_time(sample_id).unwrap(); // XXX
624        let rendering_offset = self.sample_rendering_offset(sample_id);
625        let is_sync = self.is_sync_sample(sample_id);
626
627        Ok(Some(Mp4Sample {
628            start_time,
629            duration,
630            rendering_offset,
631            is_sync,
632            bytes: Bytes::from(buffer),
633        }))
634    }
635}
636
637// TODO creation_time, modification_time
638#[derive(Debug, Default)]
639pub(crate) struct Mp4TrackWriter {
640    trak: TrakBox,
641
642    sample_id: u32,
643    fixed_sample_size: u32,
644    is_fixed_sample_size: bool,
645    chunk_samples: u32,
646    chunk_duration: u32,
647    chunk_buffer: BytesMut,
648
649    samples_per_chunk: u32,
650    duration_per_chunk: u32,
651}
652
653impl Mp4TrackWriter {
654    pub(crate) fn new(track_id: u32, config: &TrackConfig) -> Result<Self> {
655        let mut trak = TrakBox::default();
656        trak.tkhd.track_id = track_id;
657        trak.mdia.mdhd.timescale = config.timescale;
658        trak.mdia.mdhd.language = config.language.to_owned();
659        trak.mdia.hdlr.handler_type = config.track_type.into();
660        trak.mdia.minf.stbl.co64 = Some(Co64Box::default());
661        match config.media_conf {
662            MediaConfig::AvcConfig(ref avc_config) => {
663                trak.tkhd.set_width(avc_config.width);
664                trak.tkhd.set_height(avc_config.height);
665
666                let vmhd = VmhdBox::default();
667                trak.mdia.minf.vmhd = Some(vmhd);
668
669                let avc1 = Avc1Box::new(avc_config);
670                trak.mdia.minf.stbl.stsd.avc1 = Some(avc1);
671            }
672            MediaConfig::HevcConfig(ref hevc_config) => {
673                trak.tkhd.set_width(hevc_config.width);
674                trak.tkhd.set_height(hevc_config.height);
675
676                let vmhd = VmhdBox::default();
677                trak.mdia.minf.vmhd = Some(vmhd);
678
679                let hev1 = Hev1Box::new(hevc_config);
680                trak.mdia.minf.stbl.stsd.hev1 = Some(hev1);
681            }
682            MediaConfig::Vp9Config(ref config) => {
683                trak.tkhd.set_width(config.width);
684                trak.tkhd.set_height(config.height);
685
686                trak.mdia.minf.stbl.stsd.vp09 = Some(Vp09Box::new(config));
687            }
688            MediaConfig::AacConfig(ref aac_config) => {
689                let smhd = SmhdBox::default();
690                trak.mdia.minf.smhd = Some(smhd);
691
692                let mp4a = Mp4aBox::new(aac_config);
693                trak.mdia.minf.stbl.stsd.mp4a = Some(mp4a);
694            }
695            MediaConfig::TtxtConfig(ref _ttxt_config) => {
696                let tx3g = Tx3gBox::default();
697                trak.mdia.minf.stbl.stsd.tx3g = Some(tx3g);
698            }
699        }
700        Ok(Mp4TrackWriter {
701            trak,
702            chunk_buffer: BytesMut::new(),
703            sample_id: 1,
704            duration_per_chunk: config.timescale, // 1 second
705            ..Self::default()
706        })
707    }
708
709    fn update_sample_sizes(&mut self, size: u32) {
710        if self.trak.mdia.minf.stbl.stsz.sample_count == 0 {
711            if size == 0 {
712                self.trak.mdia.minf.stbl.stsz.sample_size = 0;
713                self.is_fixed_sample_size = false;
714                self.trak.mdia.minf.stbl.stsz.sample_sizes.push(0);
715            } else {
716                self.trak.mdia.minf.stbl.stsz.sample_size = size;
717                self.fixed_sample_size = size;
718                self.is_fixed_sample_size = true;
719            }
720        } else if self.is_fixed_sample_size {
721            if self.fixed_sample_size != size {
722                self.is_fixed_sample_size = false;
723                if self.trak.mdia.minf.stbl.stsz.sample_size > 0 {
724                    self.trak.mdia.minf.stbl.stsz.sample_size = 0;
725                    for _ in 0..self.trak.mdia.minf.stbl.stsz.sample_count {
726                        self.trak
727                            .mdia
728                            .minf
729                            .stbl
730                            .stsz
731                            .sample_sizes
732                            .push(self.fixed_sample_size);
733                    }
734                }
735                self.trak.mdia.minf.stbl.stsz.sample_sizes.push(size);
736            }
737        } else {
738            self.trak.mdia.minf.stbl.stsz.sample_sizes.push(size);
739        }
740        self.trak.mdia.minf.stbl.stsz.sample_count += 1;
741    }
742
743    fn update_sample_times(&mut self, dur: u32) {
744        if let Some(ref mut entry) = self.trak.mdia.minf.stbl.stts.entries.last_mut() {
745            if entry.sample_delta == dur {
746                entry.sample_count += 1;
747                return;
748            }
749        }
750
751        let entry = SttsEntry {
752            sample_count: 1,
753            sample_delta: dur,
754        };
755        self.trak.mdia.minf.stbl.stts.entries.push(entry);
756    }
757
758    fn update_rendering_offsets(&mut self, offset: i32) {
759        let ctts = if let Some(ref mut ctts) = self.trak.mdia.minf.stbl.ctts {
760            ctts
761        } else {
762            if offset == 0 {
763                return;
764            }
765            let mut ctts = CttsBox::default();
766            if self.sample_id > 1 {
767                let entry = CttsEntry {
768                    sample_count: self.sample_id - 1,
769                    sample_offset: 0,
770                };
771                ctts.entries.push(entry);
772            }
773            self.trak.mdia.minf.stbl.ctts = Some(ctts);
774            self.trak.mdia.minf.stbl.ctts.as_mut().unwrap()
775        };
776
777        if let Some(ref mut entry) = ctts.entries.last_mut() {
778            if entry.sample_offset == offset {
779                entry.sample_count += 1;
780                return;
781            }
782        }
783
784        let entry = CttsEntry {
785            sample_count: 1,
786            sample_offset: offset,
787        };
788        ctts.entries.push(entry);
789    }
790
791    fn update_sync_samples(&mut self, is_sync: bool) {
792        if let Some(ref mut stss) = self.trak.mdia.minf.stbl.stss {
793            if !is_sync {
794                return;
795            }
796
797            stss.entries.push(self.sample_id);
798        } else {
799            if !is_sync {
800                return;
801            }
802
803            // Create the stts box if not found and push the entry.
804            let mut stss = StssBox::default();
805            stss.entries.push(self.sample_id);
806            self.trak.mdia.minf.stbl.stss = Some(stss);
807        };
808    }
809
810    fn is_chunk_full(&self) -> bool {
811        if self.samples_per_chunk > 0 {
812            self.chunk_samples >= self.samples_per_chunk
813        } else {
814            self.chunk_duration >= self.duration_per_chunk
815        }
816    }
817
818    fn update_durations(&mut self, dur: u32, movie_timescale: u32) {
819        self.trak.mdia.mdhd.duration += dur as u64;
820        if self.trak.mdia.mdhd.duration > (u32::MAX as u64) {
821            self.trak.mdia.mdhd.version = 1
822        }
823        self.trak.tkhd.duration +=
824            dur as u64 * movie_timescale as u64 / self.trak.mdia.mdhd.timescale as u64;
825        if self.trak.tkhd.duration > (u32::MAX as u64) {
826            self.trak.tkhd.version = 1
827        }
828    }
829
830    pub(crate) fn write_sample<W: Write + Seek>(
831        &mut self,
832        writer: &mut W,
833        sample: &Mp4Sample,
834        movie_timescale: u32,
835    ) -> Result<u64> {
836        self.chunk_buffer.extend_from_slice(&sample.bytes);
837        self.chunk_samples += 1;
838        self.chunk_duration += sample.duration;
839        self.update_sample_sizes(sample.bytes.len() as u32);
840        self.update_sample_times(sample.duration);
841        self.update_rendering_offsets(sample.rendering_offset);
842        self.update_sync_samples(sample.is_sync);
843        if self.is_chunk_full() {
844            self.write_chunk(writer)?;
845        }
846        self.update_durations(sample.duration, movie_timescale);
847
848        self.sample_id += 1;
849
850        Ok(self.trak.tkhd.duration)
851    }
852
853    fn chunk_count(&self) -> u32 {
854        let co64 = self.trak.mdia.minf.stbl.co64.as_ref().unwrap();
855        co64.entries.len() as u32
856    }
857
858    fn update_sample_to_chunk(&mut self, chunk_id: u32) {
859        if let Some(entry) = self.trak.mdia.minf.stbl.stsc.entries.last() {
860            if entry.samples_per_chunk == self.chunk_samples {
861                return;
862            }
863        }
864
865        let entry = StscEntry {
866            first_chunk: chunk_id,
867            samples_per_chunk: self.chunk_samples,
868            sample_description_index: 1,
869            first_sample: self.sample_id - self.chunk_samples + 1,
870        };
871        self.trak.mdia.minf.stbl.stsc.entries.push(entry);
872    }
873
874    fn update_chunk_offsets(&mut self, offset: u64) {
875        let co64 = self.trak.mdia.minf.stbl.co64.as_mut().unwrap();
876        co64.entries.push(offset);
877    }
878
879    fn write_chunk<W: Write + Seek>(&mut self, writer: &mut W) -> Result<()> {
880        if self.chunk_buffer.is_empty() {
881            return Ok(());
882        }
883        let chunk_offset = writer.stream_position()?;
884
885        writer.write_all(&self.chunk_buffer)?;
886
887        self.update_sample_to_chunk(self.chunk_count() + 1);
888        self.update_chunk_offsets(chunk_offset);
889
890        self.chunk_buffer.clear();
891        self.chunk_samples = 0;
892        self.chunk_duration = 0;
893
894        Ok(())
895    }
896
897    fn max_sample_size(&self) -> u32 {
898        if self.trak.mdia.minf.stbl.stsz.sample_size > 0 {
899            self.trak.mdia.minf.stbl.stsz.sample_size
900        } else {
901            let mut max_size = 0;
902            for sample_size in self.trak.mdia.minf.stbl.stsz.sample_sizes.iter() {
903                max_size = cmp::max(max_size, *sample_size);
904            }
905            max_size
906        }
907    }
908
909    pub(crate) fn write_end<W: Write + Seek>(&mut self, writer: &mut W) -> Result<TrakBox> {
910        self.write_chunk(writer)?;
911
912        let max_sample_size = self.max_sample_size();
913        if let Some(ref mut mp4a) = self.trak.mdia.minf.stbl.stsd.mp4a {
914            if let Some(ref mut esds) = mp4a.esds {
915                esds.es_desc.dec_config.buffer_size_db = max_sample_size;
916            }
917            // TODO
918            // mp4a.esds.es_desc.dec_config.max_bitrate
919            // mp4a.esds.es_desc.dec_config.avg_bitrate
920        }
921        if let Ok(stco) = StcoBox::try_from(self.trak.mdia.minf.stbl.co64.as_ref().unwrap()) {
922            self.trak.mdia.minf.stbl.stco = Some(stco);
923            self.trak.mdia.minf.stbl.co64 = None;
924        }
925
926        Ok(self.trak.clone())
927    }
928}
mp4/track.rs

mp4/
track.rs
