Make it more robust

gyroflow · Aug 23, 2022 · ebcf70f · ebcf70f
1 parent 775a479
commit ebcf70f
Show file tree

Hide file tree

Showing 7 changed files with 58 additions and 30 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,11 +1,12 @@
 [package]
 name = "mp4-merge"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2021"
 authors = ["Adrian <[email protected]>"]
 license = "MIT OR Apache-2.0"
 description = "A tool and library to losslessly join multiple .mp4 files shot with same camera and settings"
 repository = "https://github.com/gyroflow/mp4-merge"
+readme = "README.md"
 
 [dependencies]
 byteorder = "1.4.3"

diff --git a/README.md b/README.md
@@ -10,7 +10,7 @@ All original tracks are preserved, all metadata is kept as in the original.
 It was created to help stabilizing such files in [Gyroflow](https://github.com/gyroflow/gyroflow).
 
 ## Download:
-See the [Releases](https://github.com/gyroflow/mp4_parse/releases) page.
+See the [Releases](https://github.com/gyroflow/mp4-merge/releases) page.
 
 ## Usage:
 
@@ -28,7 +28,7 @@ mp4_merge IN_FILE1.mp4 IN_FILE2.mp4 IN_FILE3.mp4 ... --out result.mp4
 
 ```toml
 [dependencies]
-mp4_merge = "0.1.0"
+mp4-merge = "0.1.1"
 ```
 ```rust
 let files = ["IN_FILE1.mp4", "IN_FILE2.mp4"];

diff --git a/src/bin.rs b/src/bin.rs
@@ -27,7 +27,7 @@ fn main() {
         println!("Merging file {:?}", p);
         files.push(p.to_owned());
         if output_file.is_none() {
-            output_file = Some(p.with_file_name(format!("{}_joined.mp4", p.to_str().unwrap())));
+            output_file = Some(p.with_file_name(format!("{}_joined.mp4", p.file_name().unwrap().to_str().unwrap())));
         }
     }
     if files.is_empty() { eprintln!("No input files!"); return; }

diff --git a/src/desc_reader.rs b/src/desc_reader.rs
@@ -15,7 +15,10 @@ pub struct TrackDesc {
     pub stsz: Vec<u32>,
     pub stco: Vec<u64>,
     pub stss: Vec<u32>,
+    pub sdtp: Vec<u8>,
     pub stss_offset: u32,
+    pub stsz_sample_size: u32,
+    pub stsz_count: u32,
 }
 
 #[derive(Default, Clone, Debug)]
@@ -24,6 +27,7 @@ pub struct Desc {
     pub moov_mvhd_duration: u64,
     pub moov_tracks: Vec<TrackDesc>,
     pub mdat_offset: u64,
+    pub mdat_final_position: u64,
 }
 
 pub fn read_desc<R: Read + Seek>(d: &mut R, desc: &mut Desc, track: usize, max_read: u64) -> Result<()> {
@@ -41,7 +45,8 @@ pub fn read_desc<R: Read + Seek>(d: &mut R, desc: &mut Desc, track: usize, max_r
             log::debug!("Reading {}, offset: {}, size: {size}", typ_to_str(typ), offs);
             let org_pos = d.stream_position()?;
             if typ == fourcc("mdat") {
-                desc.mdat_position.push((None, d.stream_position()?, size - header_size as u64));
+                desc.mdat_position.push((None, org_pos, size - header_size as u64));
+                desc.mdat_final_position = org_pos;
             }
             if typ == fourcc("mvhd") || typ == fourcc("tkhd") || typ == fourcc("mdhd") {
                 let (v, _flags) = (d.read_u8()?, d.read_u24::<BigEndian>()?);
@@ -60,23 +65,36 @@ pub fn read_desc<R: Read + Seek>(d: &mut R, desc: &mut Desc, track: usize, max_r
                     }
                 }
             }
-            if typ == fourcc("elst") || typ == fourcc("stts") || typ == fourcc("stsz") || typ == fourcc("stss") || typ == fourcc("stco") || typ == fourcc("co64") {
+            if typ == fourcc("elst") || typ == fourcc("stts") || typ == fourcc("stsz") || typ == fourcc("stss") || typ == fourcc("stco") || typ == fourcc("co64") || typ == fourcc("sdtp") {
                 let track_desc = desc.moov_tracks.get_mut(tl_track).unwrap();
                 let (v, _flags) = (d.read_u8()?, d.read_u24::<BigEndian>()?);
 
-                if typ == fourcc("elst") || typ == fourcc("stsz") {
+                if typ == fourcc("elst") {
                     d.seek(SeekFrom::Current(4))?; // Skip fields
                 }
                 if typ == fourcc("elst")  {
                     track_desc.elst_segment_duration += if v == 1 { d.read_u64::<BigEndian>()? } else { d.read_u32::<BigEndian>()? as u64 };
                 }
-                if typ == fourcc("stsz") || typ == fourcc("stss") || typ == fourcc("stco") || typ == fourcc("co64") || typ == fourcc("stts") {
+                if typ == fourcc("stsz") {
+                    track_desc.stsz_sample_size = d.read_u32::<BigEndian>()?;
                     let count = d.read_u32::<BigEndian>()?;
+                    if track_desc.stsz_sample_size == 0 {
+                        for _ in 0..count { track_desc.stsz.push(d.read_u32::<BigEndian>()?); }
+                    }
+                    track_desc.stsz_count += count;
+                }
+                if typ == fourcc("sdtp") {
+                    let count = size - header_size as u64 - 4;
+                    for _ in 0..count { track_desc.sdtp.push(d.read_u8()?); }
+                }
+                if typ == fourcc("stss") || typ == fourcc("stco") || typ == fourcc("co64") || typ == fourcc("stts") {
+                    let count = d.read_u32::<BigEndian>()?;
+                    let current_file_mdat_position = desc.mdat_position.last().unwrap().1;
+                    let mdat_offset = desc.mdat_offset as i64 - current_file_mdat_position as i64;
                     for _ in 0..count {
-                        if typ == fourcc("stsz") { track_desc.stsz.push(d.read_u32::<BigEndian>()?); }
                         if typ == fourcc("stss") { track_desc.stss.push(d.read_u32::<BigEndian>()? + track_desc.stss_offset); }
-                        if typ == fourcc("stco") { track_desc.stco.push(d.read_u32::<BigEndian>()? as u64 + desc.mdat_offset); }
-                        if typ == fourcc("co64") { track_desc.stco.push(d.read_u64::<BigEndian>()? + desc.mdat_offset); }
+                        if typ == fourcc("stco") { track_desc.stco.push((d.read_u32::<BigEndian>()? as i64 + mdat_offset) as u64); }
+                        if typ == fourcc("co64") { track_desc.stco.push((d.read_u64::<BigEndian>()? as i64 + mdat_offset) as u64); }
                         if typ == fourcc("stts") { track_desc.stts.push((d.read_u32::<BigEndian>()?, d.read_u32::<BigEndian>()?)); }
                     }
                 }

diff --git a/src/lib.rs b/src/lib.rs
@@ -3,6 +3,7 @@
 
 use std::io::{ Read, Seek, Result };
 use std::path::{ Path, PathBuf };
+use byteorder::{ ReadBytesExt, BigEndian };
 use std::time::Instant;
 
 mod desc_reader;
@@ -36,15 +37,10 @@ fn typ_to_str(typ: u32) -> String {
 
 pub fn read_box<R: Read + Seek>(reader: &mut R) -> Result<(u32, u64, u64, i64)> {
     let pos = reader.stream_position()?;
-    let mut buf = [0u8; 8];
-    reader.read_exact(&mut buf)?;
-
-    let size = u32::from_be_bytes(buf[0..4].try_into().unwrap());
-    let typ = u32::from_be_bytes(buf[4..8].try_into().unwrap());
-
+    let size = reader.read_u32::<BigEndian>()?;
+    let typ = reader.read_u32::<BigEndian>()?;
     if size == 1 {
-        reader.read_exact(&mut buf)?;
-        let largesize = u64::from_be_bytes(buf);
+        let largesize = reader.read_u64::<BigEndian>()?;
         Ok((typ, pos, largesize - 8, 16))
     } else {
         Ok((typ, pos, size as u64, 8))
@@ -66,7 +62,7 @@ pub fn join_files<P: AsRef<Path> + AsRef<std::ffi::OsStr>, F: Fn(f64)>(files: &[
             mdat.0 = Some(PathBuf::from(path));
             desc.mdat_offset += mdat.2;
             for t in &mut desc.moov_tracks {
-                t.stss_offset += t.stsz.len() as u32;
+                t.stss_offset = t.stsz_count;
             }
         }
 
@@ -83,7 +79,7 @@ pub fn join_files<P: AsRef<Path> + AsRef<std::ffi::OsStr>, F: Fn(f64)>(files: &[
             debounce = Instant::now();
         }
     });
-    writer::rewrite_from_desc(&mut f1, &mut f_out, &desc, 0, u64::MAX).unwrap();
+    writer::rewrite_from_desc(&mut f1, &mut f_out, &mut desc, 0, u64::MAX).unwrap();
     progress_cb(1.0);
 
     Ok(())

diff --git a/src/writer.rs b/src/writer.rs
@@ -5,7 +5,7 @@ use std::io::{ Read, Write, Seek, Result, SeekFrom };
 use byteorder::{ ReadBytesExt, WriteBytesExt, BigEndian };
 use crate::{ fourcc, read_box, typ_to_str, desc_reader::Desc };
 
-pub fn rewrite_from_desc<R: Read + Seek, W: Write + Seek>(d: &mut R, output_file: &mut W, desc: &Desc, track: usize, max_read: u64) -> Result<u64> {
+pub fn rewrite_from_desc<R: Read + Seek, W: Write + Seek>(d: &mut R, output_file: &mut W, desc: &mut Desc, track: usize, max_read: u64) -> Result<u64> {
     let mut total_read_size = 0;
     let mut total_new_size = 0;
     let mut tl_track = track;
@@ -37,6 +37,8 @@ pub fn rewrite_from_desc<R: Read + Seek, W: Write + Seek>(d: &mut R, output_file
             output_file.write(&0u64.to_be_bytes())?;
             new_size = 16;
 
+            desc.mdat_final_position = output_file.stream_position()?;
+
             // Merge all mdats
             for (fpath, mo, ms) in &desc.mdat_position {
                 if let Some(fpath) = fpath {
@@ -80,7 +82,7 @@ pub fn rewrite_from_desc<R: Read + Seek, W: Write + Seek>(d: &mut R, output_file
                 }
             }
 
-        } else if typ == fourcc("stts") || typ == fourcc("stsz") || typ == fourcc("stss") || typ == fourcc("stco") || typ == fourcc("co64") {
+        } else if typ == fourcc("stts") || typ == fourcc("stsz") || typ == fourcc("stss") || typ == fourcc("stco") || typ == fourcc("co64") || typ == fourcc("sdtp") {
             log::debug!("Writing new {}, offset: {}, size: {size}", typ_to_str(typ), offs);
 
             d.seek(SeekFrom::Current(size as i64 - header_size))?;
@@ -94,17 +96,26 @@ pub fn rewrite_from_desc<R: Read + Seek, W: Write + Seek>(d: &mut R, output_file
 
             let track_desc = desc.moov_tracks.get(tl_track).unwrap();
             if typ == fourcc("stts") {
-                output_file.write_u32::<BigEndian>(track_desc.stts.len() as u32)?;
+                let mut new_stts: Vec<(u32, u32)> = Vec::with_capacity(track_desc.stts.len());
+                let mut prev_delta = None;
+                for x in &track_desc.stts {
+                    if let Some(prev_delta) = prev_delta {
+                        if prev_delta == x.1 { (*new_stts.last_mut().unwrap()).0 += x.0; continue; }
+                    }
+                    prev_delta = Some(x.1);
+                    new_stts.push(*x);
+                }
+                output_file.write_u32::<BigEndian>(new_stts.len() as u32)?;
                 new_size += 4;
-                for (count, delta) in &track_desc.stts {
+                for (count, delta) in &new_stts {
                     output_file.write_u32::<BigEndian>(*count)?;
                     output_file.write_u32::<BigEndian>(*delta)?;
                     new_size += 8;
                 }
             }
             if typ == fourcc("stsz") {
-                output_file.write_u32::<BigEndian>(0)?; // sample_size
-                output_file.write_u32::<BigEndian>(track_desc.stsz.len() as u32)?;
+                output_file.write_u32::<BigEndian>(track_desc.stsz_sample_size)?; // sample_size
+                output_file.write_u32::<BigEndian>(track_desc.stsz_count)?;
                 new_size += 8;
                 for x in &track_desc.stsz { output_file.write_u32::<BigEndian>(*x as u32)?; new_size += 4; }
             }
@@ -117,12 +128,14 @@ pub fn rewrite_from_desc<R: Read + Seek, W: Write + Seek>(d: &mut R, output_file
                 output_file.write_u32::<BigEndian>(track_desc.stco.len() as u32)?;
                 new_size += 4;
                 for x in &track_desc.stco {
-                    output_file.write_u64::<BigEndian>(*x + 8)?; // TODO: + 8 only if the original mdat was not large box already
+                    output_file.write_u64::<BigEndian>(*x + desc.mdat_final_position)?;
                     new_size += 8;
                 }
             }
+            if typ == fourcc("sdtp") {
+                for x in &track_desc.sdtp { output_file.write_u8(*x)?; new_size += 1; }
+            }
             patch_bytes(output_file, out_pos, &(new_size as u32).to_be_bytes())?;
-
         } else {
             log::debug!("Writing original {}, offset: {}, size: {size}", typ_to_str(typ), offs);