more boxes

3 files changed, 262 insertions, 49 deletions

diff --git a/remuxer/mp4/src/boxes.rs b/remuxer/mp4/src/boxes.rs
index 4c14eb1..f695d98 100644
--- a/remuxer/mp4/src/boxes.rs
+++ b/remuxer/mp4/src/boxes.rs
@@ -27,6 +27,7 @@ container_box!(DataInformation, b"dinf");
 container_box!(SampleTable, b"stbl");
 container_box!(MovieFragment, b"stbl");
 container_box!(TrackFragment, b"stbl");
+container_box!(MovieExtends, b"mvex");
 
 pub struct FileType<'a> {
     pub major_brand: [u8; 4],
@@ -106,7 +107,6 @@ impl Default for MovieHeader {
     }
 }
 
-#[derive(Default)]
 pub struct TrackHeader {
     pub creation_time: u64,
     pub modification_time: u64,
@@ -140,10 +140,24 @@ impl WriteBox for TrackHeader {
         buf.extend(self.height.to_be_bytes());
     }
 }
-
-pub struct EditList<'a> {
-    pub entries: &'a [EditListEntry],
+impl Default for TrackHeader {
+    fn default() -> Self {
+        Self {
+            creation_time: 0,
+            modification_time: 0,
+            track_id: 0,
+            duration: 0,
+            layer: 0,
+            alternate_group: 0,
+            volume: 0,
+            matrix: [0x00010000, 0, 0, 0, 0x00010000, 0, 0, 0, 0x00010000],
+            width: 0,
+            height: 0,
+        }
+    }
 }
+
+pub struct EditList<'a>(pub &'a [EditListEntry]);
 pub struct EditListEntry {
     pub edit_duration: u64,
     pub media_time: u64,
@@ -153,8 +167,8 @@ impl WriteBox for EditList<'_> {
     const BOXTYPE: [u8; 4] = *b"elst";
     const VERSION: Option<u8> = Some(1);
     fn write(&self, buf: &mut Vec<u8>) {
-        buf.extend((self.entries.len() as u32).to_be_bytes());
-        for e in self.entries {
+        buf.extend((self.0.len() as u32).to_be_bytes());
+        for e in self.0 {
             buf.extend(e.edit_duration.to_be_bytes());
             buf.extend(e.media_time.to_be_bytes());
             buf.extend(e.media_rate.to_be_bytes());
@@ -189,27 +203,28 @@ impl WriteBox for MediaHeader {
     }
 }
 
-pub struct Handler {
-    pub handler_type: u32,
-    pub name: String,
+pub struct Handler<'a> {
+    pub handler_type: [u8; 4],
+    pub name: &'a str,
 }
-impl WriteBox for Handler {
+impl WriteBox for Handler<'_> {
     const BOXTYPE: [u8; 4] = *b"hdlr";
     const VERSION: Option<u8> = Some(0);
     fn write(&self, buf: &mut Vec<u8>) {
         buf.extend([0; 4]);
-        buf.extend(self.handler_type.to_be_bytes());
+        buf.extend(self.handler_type);
         buf.extend([0; 12]);
         buf.extend(self.name.as_bytes());
+        buf.push(0);
     }
 }
 
 #[derive(Default)]
-pub struct VideoMediaHandler {
+pub struct VideoMediaHeader {
     pub graphicsmode: u16,
     pub opcolor: [u16; 3],
 }
-impl WriteBox for VideoMediaHandler {
+impl WriteBox for VideoMediaHeader {
     const BOXTYPE: [u8; 4] = *b"vmhd";
     const VERSION: Option<u8> = Some(0);
     fn write(&self, buf: &mut Vec<u8>) {
@@ -221,10 +236,10 @@ impl WriteBox for VideoMediaHandler {
 }
 
 #[derive(Default)]
-pub struct SoundMediaHandler {
+pub struct SoundMediaHeader {
     pub balance: u16,
 }
-impl WriteBox for SoundMediaHandler {
+impl WriteBox for SoundMediaHeader {
     const BOXTYPE: [u8; 4] = *b"smhd";
     const VERSION: Option<u8> = Some(0);
     fn write(&self, buf: &mut Vec<u8>) {
@@ -232,14 +247,29 @@ impl WriteBox for SoundMediaHandler {
     }
 }
 
-pub struct DataReference<'a> {
-    pub entries: &'a [()],
+pub struct DataReference<'a>(pub &'a [DataReferenceEntry]);
+pub enum DataReferenceEntry {
+    URL(String),
+    URN { name: String, location: String },
+    Imda(u32),
+    SeqNumImda,
 }
 impl WriteBox for DataReference<'_> {
     const BOXTYPE: [u8; 4] = *b"dref";
     const VERSION: Option<u8> = Some(0);
     fn write(&self, buf: &mut Vec<u8>) {
-        buf.extend([0; 4]);
+        buf.extend((self.0.len() as u32).to_be_bytes());
+        for e in self.0 {
+            match e {
+                DataReferenceEntry::URL(x) if x.is_empty() => {
+                    buf.extend(12u32.to_be_bytes());
+                    buf.extend(b"url ");
+                    buf.push(0); // version
+                    buf.extend([0, 0, 1]); // flags
+                }
+                _ => todo!(),
+            }
+        }
     }
 }
 
@@ -353,9 +383,130 @@ impl WriteBox for ChunkOffset<'_> {
     }
 }
 
-pub struct SampleDescription;
-impl WriteBox for SampleDescription {
+pub struct SyncSample<'a>(pub &'a [u32]);
+impl WriteBox for SyncSample<'_> {
+    const BOXTYPE: [u8; 4] = *b"stss";
+    const VERSION: Option<u8> = Some(0);
+    fn write(&self, buf: &mut Vec<u8>) {
+        buf.extend((self.0.len() as u32).to_be_bytes());
+        for e in self.0 {
+            buf.extend(e.to_be_bytes());
+        }
+    }
+}
+
+pub struct SampleDescription<T: Fn(&mut BoxW)>(pub T);
+impl<T: Fn(&mut BoxW)> WriteBox for SampleDescription<T> {
     const BOXTYPE: [u8; 4] = *b"stsd";
     const VERSION: Option<u8> = Some(0);
-    fn write(&self, _buf: &mut Vec<u8>) {}
+    fn write(&self, buf: &mut Vec<u8>) {
+        buf.extend(1u32.to_be_bytes());
+        self.0(&mut BoxW(buf))
+    }
+}
+
+pub struct AVCSampleEntry<'a> {
+    pub visual: VisualSampleEntry,
+    pub config: AVCConfiguration<'a>,
+}
+impl WriteBox for AVCSampleEntry<'_> {
+    const BOXTYPE: [u8; 4] = *b"avc1";
+    fn write(&self, buf: &mut Vec<u8>) {
+        self.visual.write(buf);
+        BoxW(buf).write(&self.config);
+    }
+}
+
+pub struct AVCConfiguration<'a>(pub &'a [u8]);
+impl WriteBox for AVCConfiguration<'_> {
+    const BOXTYPE: [u8; 4] = *b"avcC";
+    fn write(&self, buf: &mut Vec<u8>) {
+        buf.extend(self.0);
+    }
+}
+
+#[derive(Default)]
+pub struct SampleEntry {
+    pub data_reference_index: u16,
+}
+impl WriteBox for SampleEntry {
+    const BOXTYPE: [u8; 4] = [0; 4];
+    fn write(&self, buf: &mut Vec<u8>) {
+        buf.extend([0; 6]);
+        buf.extend(self.data_reference_index.to_be_bytes())
+    }
+}
+
+pub struct VisualSampleEntry {
+    pub sample: SampleEntry,
+    pub width: u16,
+    pub height: u16,
+    pub horizresolution: u32,
+    pub vertresolution: u32,
+    pub frame_count: u16,
+    pub compressorname: [u8; 32],
+    pub depth: u16,
+}
+impl WriteBox for VisualSampleEntry {
+    const BOXTYPE: [u8; 4] = [0; 4];
+    fn write(&self, buf: &mut Vec<u8>) {
+        self.sample.write(buf);
+        buf.extend([0; 2]);
+        buf.extend([0; 2]);
+        buf.extend([0; 12]);
+        buf.extend(self.width.to_be_bytes());
+        buf.extend(self.height.to_be_bytes());
+        buf.extend(self.horizresolution.to_be_bytes());
+        buf.extend(self.vertresolution.to_be_bytes());
+        buf.extend([0; 4]);
+        buf.extend(self.frame_count.to_be_bytes());
+        buf.extend(self.compressorname);
+        buf.extend(self.depth.to_be_bytes());
+        buf.extend((-1i16).to_be_bytes());
+    }
+}
+impl Default for VisualSampleEntry {
+    fn default() -> Self {
+        Self {
+            sample: SampleEntry::default(),
+            width: 0,
+            height: 0,
+            horizresolution: 0x00480000,
+            vertresolution: 0x00480000,
+            frame_count: 1,
+            compressorname: [0; 32],
+            depth: 0x0018,
+        }
+    }
+}
+
+pub struct MovieExtendsHeader {
+    pub fragment_duration: u64,
+}
+impl WriteBox for MovieExtendsHeader {
+    const BOXTYPE: [u8; 4] = *b"mehd";
+    const VERSION: Option<u8> = Some(1);
+    fn write(&self, buf: &mut Vec<u8>) {
+        buf.extend(self.fragment_duration.to_be_bytes())
+    }
+}
+
+#[derive(Default)]
+pub struct TrackExtends {
+    pub track_id: u32,
+    pub default_sample_description_index: u32,
+    pub default_sample_duration: u32,
+    pub default_sample_size: u32,
+    pub default_sample_flags: u32,
+}
+impl WriteBox for TrackExtends {
+    const BOXTYPE: [u8; 4] = *b"trex";
+    const VERSION: Option<u8> = Some(0);
+    fn write(&self, buf: &mut Vec<u8>) {
+        buf.extend(self.track_id.to_be_bytes());
+        buf.extend(self.default_sample_description_index.to_be_bytes());
+        buf.extend(self.default_sample_duration.to_be_bytes());
+        buf.extend(self.default_sample_size.to_be_bytes());
+        buf.extend(self.default_sample_flags.to_be_bytes());
+    }
 }
diff --git a/remuxer/mp4/src/lib.rs b/remuxer/mp4/src/lib.rs
index cb728fa..7af8417 100644
--- a/remuxer/mp4/src/lib.rs
+++ b/remuxer/mp4/src/lib.rs
@@ -35,3 +35,14 @@ pub trait WriteBox {
         0
     }
 }
+
+impl<T: WriteBox> WriteBox for &T {
+    const BOXTYPE: [u8; 4] = T::BOXTYPE;
+    const VERSION: Option<u8> = T::VERSION;
+    fn write(&self, buf: &mut Vec<u8>) {
+        (**self).write(buf);
+    }
+    fn flags(&self) -> u32 {
+        (**self).flags()
+    }
+}
diff --git a/remuxer/src/muxers/mp4.rs b/remuxer/src/muxers/mp4.rs
index 5975921..bd36789 100644
--- a/remuxer/src/muxers/mp4.rs
+++ b/remuxer/src/muxers/mp4.rs
@@ -6,7 +6,7 @@
 
 use crate::muxers::FragmentMuxer;
 use anyhow::Result;
-use winter_matroska::Segment;
+use winter_matroska::{Segment, TrackType};
 use winter_mp4::*;
 
 pub struct MP4FragmentMuxer;
@@ -18,42 +18,93 @@ impl FragmentMuxer for MP4FragmentMuxer {
             minor_version: 512,
             compatible_brands: &[*b"iso5", *b"iso6", *b"mp41"],
         });
+
         out.write(Movie(|moov| {
             moov.write(MovieHeader {
+                timescale: 1_000_000_000,
                 ..Default::default()
             });
-            moov.write(Track(|trak| {
-                trak.write(TrackHeader {
-                    ..Default::default()
-                });
-                trak.write(Edit(|edts| {
-                    edts.write(EditList { entries: &[] });
-                }));
-                trak.write(Media(|mdia| {
-                    mdia.write(Handler {
-                        handler_type: 0,
-                        name: "blub".to_string(),
+            for track in &segment.tracks.as_ref().unwrap().entries {
+                moov.write(Track(|trak| {
+                    trak.write(TrackHeader {
+                        width: track.video.as_ref().map_or(0, |v| v.pixel_width as u32),
+                        height: track.video.as_ref().map_or(0, |v| v.pixel_height as u32),
+                        volume: track.audio.as_ref().map_or(0, |_| 0x0100),
+                        track_id: track.track_number as u32,
+                        ..Default::default()
                     });
-                    mdia.write(MediaInformation(|minf| {
-                        minf.write(VideoMediaHandler {
-                            ..Default::default()
+                    // trak.write(Edit(|edts| {
+                    //     edts.write(EditList(&[]));
+                    // }));
+                    trak.write(Media(|mdia| {
+                        mdia.write(Handler {
+                            handler_type: match track.track_type {
+                                TrackType::Video => *b"vide",
+                                TrackType::Audio => *b"soun",
+                                _ => todo!(),
+                            },
+                            name: &format!("ISO Media file produced by {}", env!("CARGO_PKG_NAME")),
                         });
-                        minf.write(DataInformation(|dinf| {
-                            dinf.write(DataReference { entries: &[] });
-                        }));
-                        minf.write(SampleTable(|stbl| {
-                            stbl.write(SampleDescription);
-                            stbl.write(TimeToSample(&[]));
-                            stbl.write(TimeToChunk(&[]));
-                            stbl.write(ChunkOffset(&[]));
-                            stbl.write(SampleSize {
-                                sample_size: 0,
-                                entries: &[],
-                            });
+                        mdia.write(MediaInformation(|minf| {
+                            match track.track_type {
+                                TrackType::Video => minf.write(VideoMediaHeader::default()),
+                                TrackType::Audio => minf.write(SoundMediaHeader::default()),
+                                _ => todo!(),
+                            };
+                            minf.write(DataInformation(|dinf| {
+                                dinf.write(DataReference(&[DataReferenceEntry::URL(
+                                    String::default(),
+                                )]));
+                            }));
+                            minf.write(SampleTable(|stbl| {
+                                stbl.write(SampleDescription(|stsd| {
+                                    if let Some(video) = &track.video
+                                        && track.codec_id == "V_MPEG4/ISO/AVC"
+                                    {
+                                        stsd.write(AVCSampleEntry {
+                                            visual: VisualSampleEntry {
+                                                sample: SampleEntry {
+                                                    data_reference_index: 1,
+                                                },
+                                                width: video.pixel_width as u16,
+                                                height: video.pixel_height as u16,
+                                                ..Default::default()
+                                            },
+                                            config: AVCConfiguration(
+                                                track.codec_private.as_ref().unwrap(),
+                                            ),
+                                        });
+                                    } else {
+                                        todo!()
+                                    }
+                                }));
+                                stbl.write(TimeToSample(&[]));
+                                stbl.write(TimeToChunk(&[]));
+                                stbl.write(ChunkOffset(&[]));
+                                stbl.write(SyncSample(&[]));
+                                stbl.write(SampleSize {
+                                    sample_size: 0,
+                                    entries: &[],
+                                });
+                            }));
                         }));
                     }));
                 }));
-            }));
+                moov.write(MovieExtends(|mvex| {
+                    mvex.write(MovieExtendsHeader {
+                        fragment_duration: (segment.info.duration.unwrap_or_default()
+                            * segment.info.timestamp_scale as f64
+                            * 1e-9) as u64,
+                    });
+                    for track in &segment.tracks.as_ref().unwrap().entries {
+                        mvex.write(TrackExtends {
+                            track_id: track.track_number as u32,
+                            default_sample_description_index: 1,
+                            ..Default::default()
+                        });
+                    }
+                }));
+            }
         }));
 
         Ok(())