Fix getting samples from movie fragments (#106)

* Fix getting samples from movie fragments

* Add a function to the reader to read in fragments from a different reader

---------

Co-authored-by: Alfred Gutierrez <alfg@users.noreply.github.com>

src/reader.rs

@@ -24,6 +24,7 @@ impl<R: Read + Seek> Mp4Reader<R> {
         let mut ftyp = None;
         let mut moov = None;
         let mut moofs = Vec::new();
+        let mut moof_offsets = Vec::new();
         let mut emsgs = Vec::new();
 
         let mut current = start;
@@ -57,8 +58,10 @@ impl<R: Read + Seek> Mp4Reader<R> {
                     moov = Some(MoovBox::read_box(&mut reader, s)?);
                 }
                 BoxType::MoofBox => {
+                    let moof_offset = reader.stream_position()? - 8;
                     let moof = MoofBox::read_box(&mut reader, s)?;
                     moofs.push(moof);
+                    moof_offsets.push(moof_offset);
                 }
                 BoxType::EmsgBox => {
                     let emsg = EmsgBox::read_box(&mut reader, s)?;
@@ -101,11 +104,12 @@ impl<R: Read + Seek> Mp4Reader<R> {
                 }
             }
 
-            for moof in moofs.iter() {
+            for (moof, moof_offset) in moofs.iter().zip(moof_offsets) {
                 for traf in moof.trafs.iter() {
                     let track_id = traf.tfhd.track_id;
                     if let Some(track) = tracks.get_mut(&track_id) {
                         track.default_sample_duration = default_sample_duration;
+                        track.moof_offsets.push(moof_offset);
                         track.trafs.push(traf.clone())
                     } else {
                         return Err(Error::TrakNotFound(track_id));
@@ -125,6 +129,92 @@ impl<R: Read + Seek> Mp4Reader<R> {
         })
     }
 
+    pub fn read_fragment_header<FR: Read + Seek>(
+        &self,
+        mut reader: FR,
+        size: u64,
+    ) -> Result<Mp4Reader<FR>> {
+        let start = reader.stream_position()?;
+
+        let mut moofs = Vec::new();
+        let mut moof_offsets = Vec::new();
+
+        let mut current = start;
+        while current < size {
+            // Get box header.
+            let header = BoxHeader::read(&mut reader)?;
+            let BoxHeader { name, size: s } = header;
+            if s > size {
+                return Err(Error::InvalidData(
+                    "file contains a box with a larger size than it",
+                ));
+            }
+
+            // Break if size zero BoxHeader, which can result in dead-loop.
+            if s == 0 {
+                break;
+            }
+
+            // Match and parse the atom boxes.
+            match name {
+                BoxType::MdatBox => {
+                    skip_box(&mut reader, s)?;
+                }
+                BoxType::MoofBox => {
+                    let moof_offset = reader.stream_position()? - 8;
+                    let moof = MoofBox::read_box(&mut reader, s)?;
+                    moofs.push(moof);
+                    moof_offsets.push(moof_offset);
+                }
+                _ => {
+                    // XXX warn!()
+                    skip_box(&mut reader, s)?;
+                }
+            }
+            current = reader.stream_position()?;
+        }
+
+        if moofs.is_empty() {
+            return Err(Error::BoxNotFound(BoxType::MoofBox));
+        }
+
+        let size = current - start;
+        let mut tracks: HashMap<u32, Mp4Track> = self
+            .moov
+            .traks
+            .iter()
+            .map(|trak| (trak.tkhd.track_id, Mp4Track::from(trak)))
+            .collect();
+
+        let mut default_sample_duration = 0;
+        if let Some(ref mvex) = &self.moov.mvex {
+            default_sample_duration = mvex.trex.default_sample_duration
+        }
+
+        for (moof, moof_offset) in moofs.iter().zip(moof_offsets) {
+            for traf in moof.trafs.iter() {
+                let track_id = traf.tfhd.track_id;
+                if let Some(track) = tracks.get_mut(&track_id) {
+                    track.default_sample_duration = default_sample_duration;
+                    track.moof_offsets.push(moof_offset);
+                    track.trafs.push(traf.clone())
+                } else {
+                    return Err(Error::TrakNotFound(track_id));
+                }
+            }
+        }
+
+        Ok(Mp4Reader {
+            reader,
+            ftyp: self.ftyp.clone(),
+            moov: self.moov.clone(),
+            moofs,
+            emsgs: Vec::new(),
+            tracks,
+            size,
+        })
+    }
+
     pub fn size(&self) -> u64 {
         self.size
     }

src/track.rs

@@ -6,6 +6,7 @@ use std::time::Duration;
 
 use crate::mp4box::traf::TrafBox;
 use crate::mp4box::trak::TrakBox;
+use crate::mp4box::trun::TrunBox;
 use crate::mp4box::{
     avc1::Avc1Box, co64::Co64Box, ctts::CttsBox, ctts::CttsEntry, hev1::Hev1Box, mp4a::Mp4aBox,
     smhd::SmhdBox, stco::StcoBox, stsc::StscEntry, stss::StssBox, stts::SttsEntry, tx3g::Tx3gBox,
@@ -92,6 +93,7 @@ impl From<Vp9Config> for TrackConfig {
 pub struct Mp4Track {
     pub trak: TrakBox,
     pub trafs: Vec<TrafBox>,
+    pub moof_offsets: Vec<u64>,
 
     // Fragmented Tracks Defaults.
     pub default_sample_duration: u32,
@@ -103,6 +105,7 @@ impl Mp4Track {
         Self {
             trak,
             trafs: Vec::new(),
+            moof_offsets: Vec::new(),
             default_sample_duration: 0,
         }
     }
@@ -436,8 +439,32 @@ impl Mp4Track {
 
     fn sample_offset(&self, sample_id: u32) -> Result<u64> {
         if !self.trafs.is_empty() {
-            if let Some((traf_idx, _sample_idx)) = self.find_traf_idx_and_sample_idx(sample_id) {
-                Ok(self.trafs[traf_idx].tfhd.base_data_offset.unwrap_or(0))
+            if let Some((traf_idx, sample_idx)) = self.find_traf_idx_and_sample_idx(sample_id) {
+                let mut sample_offset = self.trafs[traf_idx]
+                    .tfhd
+                    .base_data_offset
+                    .unwrap_or(self.moof_offsets[traf_idx]);
+
+                if let Some(data_offset) = self.trafs[traf_idx]
+                    .trun
+                    .as_ref()
+                    .and_then(|trun| trun.data_offset)
+                {
+                    sample_offset = sample_offset.checked_add_signed(data_offset as i64).ok_or(
+                        Error::InvalidData("attempt to calculate trun sample offset with overflow"),
+                    )?;
+                }
+
+                let first_sample_in_trun = sample_id - sample_idx as u32;
+                for i in first_sample_in_trun..sample_id {
+                    sample_offset = sample_offset
+                        .checked_add(self.sample_size(i)? as u64)
+                        .ok_or(Error::InvalidData(
+                            "attempt to calculate trun entry sample offset with overflow",
+                        ))?;
+                }
+
+                Ok(sample_offset)
             } else {
                 Err(Error::BoxInTrafNotFound(self.track_id(), BoxType::TrafBox))
             }
@@ -473,15 +500,38 @@ impl Mp4Track {
     }
 
     fn sample_time(&self, sample_id: u32) -> Result<(u64, u32)> {
-        let stts = &self.trak.mdia.minf.stbl.stts;
-
-        let mut sample_count: u32 = 1;
-        let mut elapsed = 0;
-
         if !self.trafs.is_empty() {
-            let start_time = ((sample_id - 1) * self.default_sample_duration) as u64;
-            Ok((start_time, self.default_sample_duration))
+            let mut base_start_time = 0;
+            let mut default_sample_duration = self.default_sample_duration;
+            if let Some((traf_idx, sample_idx)) = self.find_traf_idx_and_sample_idx(sample_id) {
+                let traf = &self.trafs[traf_idx];
+                if let Some(tfdt) = &traf.tfdt {
+                    base_start_time = tfdt.base_media_decode_time;
+                }
+                if let Some(duration) = traf.tfhd.default_sample_duration {
+                    default_sample_duration = duration;
+                }
+                if let Some(trun) = &traf.trun {
+                    if TrunBox::FLAG_SAMPLE_DURATION & trun.flags != 0 {
+                        let mut start_offset = 0u64;
+                        for duration in &trun.sample_durations[..sample_idx] {
+                            start_offset = start_offset.checked_add(*duration as u64).ok_or(
+                                Error::InvalidData("attempt to sum sample durations with overflow"),
+                            )?;
+                        }
+                        let duration = trun.sample_durations[sample_idx];
+                        return Ok((base_start_time + start_offset, duration));
+                    }
+                }
+            }
+            let start_offset = ((sample_id - 1) * default_sample_duration) as u64;
+            Ok((base_start_time + start_offset, default_sample_duration))
         } else {
+            let stts = &self.trak.mdia.minf.stbl.stts;
+
+            let mut sample_count: u32 = 1;
+            let mut elapsed = 0;
+
             for entry in stts.entries.iter() {
                 let new_sample_count =
                     sample_count
@@ -508,7 +558,17 @@ impl Mp4Track {
     }
 
     fn sample_rendering_offset(&self, sample_id: u32) -> i32 {
-        if let Some(ref ctts) = self.trak.mdia.minf.stbl.ctts {
+        if !self.trafs.is_empty() {
+            if let Some((traf_idx, sample_idx)) = self.find_traf_idx_and_sample_idx(sample_id) {
+                if let Some(cts) = self.trafs[traf_idx]
+                    .trun
+                    .as_ref()
+                    .and_then(|trun| trun.sample_cts.get(sample_idx))
+                {
+                    return *cts as i32;
+                }
+            }
+        } else if let Some(ref ctts) = self.trak.mdia.minf.stbl.ctts {
             if let Ok((ctts_index, _)) = self.ctts_index(sample_id) {
                 let ctts_entry = ctts.entries.get(ctts_index).unwrap();
                 return ctts_entry.sample_offset;

tests/lib.rs

@@ -176,3 +176,36 @@ fn test_read_metadata() {
     assert_eq!(poster.len(), want_poster.len());
     assert_eq!(poster, want_poster.as_slice());
 }
+
+#[test]
+fn test_read_fragments() {
+    let mp4 = get_reader("tests/samples/minimal_init.mp4");
+
+    assert_eq!(692, mp4.size());
+    assert_eq!(5, mp4.compatible_brands().len());
+
+    let sample_count = mp4.sample_count(1).unwrap();
+    assert_eq!(sample_count, 0);
+
+    let f = File::open("tests/samples/minimal_fragment.m4s").unwrap();
+    let f_size = f.metadata().unwrap().len();
+    let frag_reader = BufReader::new(f);
+
+    let mut mp4_fragment = mp4.read_fragment_header(frag_reader, f_size).unwrap();
+    let sample_count = mp4_fragment.sample_count(1).unwrap();
+    assert_eq!(sample_count, 1);
+    let sample_1_1 = mp4_fragment.read_sample(1, 1).unwrap().unwrap();
+    assert_eq!(sample_1_1.bytes.len(), 751);
+    assert_eq!(
+        sample_1_1,
+        mp4::Mp4Sample {
+            start_time: 0,
+            duration: 512,
+            rendering_offset: 0,
+            is_sync: true,
+            bytes: mp4::Bytes::from(vec![0x0u8; 751]),
+        }
+    );
+    let eos = mp4_fragment.read_sample(1, 2);
+    assert!(eos.is_err());
+}