Fix TfhdBox optional fields (#86)

* Fix TfhdBox optional fields

* clippy fixes

* Cargo fmt fix

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

src/mp4box/tfhd.rs

@@ -9,16 +9,42 @@ pub struct TfhdBox {
     pub version: u8,
     pub flags: u32,
     pub track_id: u32,
-    pub base_data_offset: u64,
+    pub base_data_offset: Option<u64>,
+    pub sample_description_index: Option<u32>,
+    pub default_sample_duration: Option<u32>,
+    pub default_sample_size: Option<u32>,
+    pub default_sample_flags: Option<u32>,
 }
 
 impl TfhdBox {
+    pub const FLAG_BASE_DATA_OFFSET: u32 = 0x01;
+    pub const FLAG_SAMPLE_DESCRIPTION_INDEX: u32 = 0x02;
+    pub const FLAG_DEFAULT_SAMPLE_DURATION: u32 = 0x08;
+    pub const FLAG_DEFAULT_SAMPLE_SIZE: u32 = 0x10;
+    pub const FLAG_DEFAULT_SAMPLE_FLAGS: u32 = 0x20;
+
     pub fn get_type(&self) -> BoxType {
         BoxType::TfhdBox
     }
 
     pub fn get_size(&self) -> u64 {
-        HEADER_SIZE + HEADER_EXT_SIZE + 4 + 8
+        let mut sum = HEADER_SIZE + HEADER_EXT_SIZE + 4;
+        if TfhdBox::FLAG_BASE_DATA_OFFSET & self.flags > 0 {
+            sum += 8;
+        }
+        if TfhdBox::FLAG_SAMPLE_DESCRIPTION_INDEX & self.flags > 0 {
+            sum += 4;
+        }
+        if TfhdBox::FLAG_DEFAULT_SAMPLE_DURATION & self.flags > 0 {
+            sum += 4;
+        }
+        if TfhdBox::FLAG_DEFAULT_SAMPLE_SIZE & self.flags > 0 {
+            sum += 4;
+        }
+        if TfhdBox::FLAG_DEFAULT_SAMPLE_FLAGS & self.flags > 0 {
+            sum += 4;
+        }
+        sum
     }
 }
 
@@ -47,7 +73,31 @@ impl<R: Read + Seek> ReadBox<&mut R> for TfhdBox {
 
         let (version, flags) = read_box_header_ext(reader)?;
         let track_id = reader.read_u32::<BigEndian>()?;
-        let base_data_offset = reader.read_u64::<BigEndian>()?;
+        let base_data_offset = if TfhdBox::FLAG_BASE_DATA_OFFSET & flags > 0 {
+            Some(reader.read_u64::<BigEndian>()?)
+        } else {
+            None
+        };
+        let sample_description_index = if TfhdBox::FLAG_SAMPLE_DESCRIPTION_INDEX & flags > 0 {
+            Some(reader.read_u32::<BigEndian>()?)
+        } else {
+            None
+        };
+        let default_sample_duration = if TfhdBox::FLAG_DEFAULT_SAMPLE_DURATION & flags > 0 {
+            Some(reader.read_u32::<BigEndian>()?)
+        } else {
+            None
+        };
+        let default_sample_size = if TfhdBox::FLAG_DEFAULT_SAMPLE_SIZE & flags > 0 {
+            Some(reader.read_u32::<BigEndian>()?)
+        } else {
+            None
+        };
+        let default_sample_flags = if TfhdBox::FLAG_DEFAULT_SAMPLE_FLAGS & flags > 0 {
+            Some(reader.read_u32::<BigEndian>()?)
+        } else {
+            None
+        };
 
         skip_bytes_to(reader, start + size)?;
 
@@ -56,6 +106,10 @@ impl<R: Read + Seek> ReadBox<&mut R> for TfhdBox {
             flags,
             track_id,
             base_data_offset,
+            sample_description_index,
+            default_sample_duration,
+            default_sample_size,
+            default_sample_flags,
         })
     }
 }
@@ -67,7 +121,21 @@ impl<W: Write> WriteBox<&mut W> for TfhdBox {
 
         write_box_header_ext(writer, self.version, self.flags)?;
         writer.write_u32::<BigEndian>(self.track_id)?;
-        writer.write_u64::<BigEndian>(self.base_data_offset)?;
+        if let Some(base_data_offset) = self.base_data_offset {
+            writer.write_u64::<BigEndian>(base_data_offset)?;
+        }
+        if let Some(sample_description_index) = self.sample_description_index {
+            writer.write_u32::<BigEndian>(sample_description_index)?;
+        }
+        if let Some(default_sample_duration) = self.default_sample_duration {
+            writer.write_u32::<BigEndian>(default_sample_duration)?;
+        }
+        if let Some(default_sample_size) = self.default_sample_size {
+            writer.write_u32::<BigEndian>(default_sample_size)?;
+        }
+        if let Some(default_sample_flags) = self.default_sample_flags {
+            writer.write_u32::<BigEndian>(default_sample_flags)?;
+        }
 
         Ok(size)
     }
@@ -85,7 +153,38 @@ mod tests {
             version: 0,
             flags: 0,
             track_id: 1,
-            base_data_offset: 0,
+            base_data_offset: None,
+            sample_description_index: None,
+            default_sample_duration: None,
+            default_sample_size: None,
+            default_sample_flags: None,
+        };
+        let mut buf = Vec::new();
+        src_box.write_box(&mut buf).unwrap();
+        assert_eq!(buf.len(), src_box.box_size() as usize);
+
+        let mut reader = Cursor::new(&buf);
+        let header = BoxHeader::read(&mut reader).unwrap();
+        assert_eq!(header.name, BoxType::TfhdBox);
+        assert_eq!(src_box.box_size(), header.size);
+
+        let dst_box = TfhdBox::read_box(&mut reader, header.size).unwrap();
+        assert_eq!(src_box, dst_box);
+    }
+
+    #[test]
+    fn test_tfhd_with_flags() {
+        let src_box = TfhdBox {
+            version: 0,
+            flags: TfhdBox::FLAG_SAMPLE_DESCRIPTION_INDEX
+                | TfhdBox::FLAG_DEFAULT_SAMPLE_DURATION
+                | TfhdBox::FLAG_DEFAULT_SAMPLE_FLAGS,
+            track_id: 1,
+            base_data_offset: None,
+            sample_description_index: Some(1),
+            default_sample_duration: Some(512),
+            default_sample_size: None,
+            default_sample_flags: Some(0x1010000),
         };
         let mut buf = Vec::new();
         src_box.write_box(&mut buf).unwrap();

src/track.rs

@@ -427,7 +427,7 @@ 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)
+                Ok(self.trafs[traf_idx].tfhd.base_data_offset.unwrap_or(0))
             } else {
                 Err(Error::BoxInTrafNotFound(self.track_id(), BoxType::TrafBox))
             }