Skip to main content

mp4_edit/atom/leaf/
stsd.rs

1use derive_more::Display;
2
3pub use crate::atom::stsd::extension::{
4    BtrtExtension, DecoderSpecificInfo, EsdsExtension, StsdExtension,
5};
6use crate::{atom::FourCC, parser::ParseAtomData, writer::SerializeAtom, ParseError};
7
8mod audio;
9mod extension;
10mod text;
11
12pub use audio::*;
13pub use text::*;
14
15pub const STSD: FourCC = FourCC::new(b"stsd");
16
17pub const SAMPLE_ENTRY_MP4A: FourCC = FourCC::new(b"mp4a"); // AAC audio
18pub const SAMPLE_ENTRY_AAVD: FourCC = FourCC::new(b"aavd"); // Audible Audio
19pub const SAMPLE_ENTRY_TEXT: FourCC = FourCC::new(b"text"); // Plain text
20
21#[derive(Debug, Clone, Display, PartialEq)]
22#[display("{}", self.as_str())]
23pub enum SampleEntryType {
24    /// AAC audio
25    Mp4a,
26    /// Audible Audio (can be treated as Mp4a)
27    Aavd,
28    /// QuickTime Text Media
29    Text,
30    /// Unknown/unsupported sample entry type
31    Unknown(FourCC),
32}
33
34impl From<FourCC> for SampleEntryType {
35    fn from(fourcc: FourCC) -> Self {
36        match fourcc {
37            SAMPLE_ENTRY_MP4A => SampleEntryType::Mp4a,
38            SAMPLE_ENTRY_AAVD => SampleEntryType::Aavd,
39            SAMPLE_ENTRY_TEXT => SampleEntryType::Text,
40            _ => SampleEntryType::Unknown(fourcc),
41        }
42    }
43}
44
45impl From<SampleEntryType> for FourCC {
46    fn from(value: SampleEntryType) -> Self {
47        match value {
48            SampleEntryType::Mp4a => SAMPLE_ENTRY_MP4A,
49            SampleEntryType::Aavd => SAMPLE_ENTRY_AAVD,
50            SampleEntryType::Text => SAMPLE_ENTRY_TEXT,
51            SampleEntryType::Unknown(bytes) => bytes,
52        }
53    }
54}
55
56impl SampleEntryType {
57    fn into_bytes(self) -> [u8; 4] {
58        FourCC::from(self).into_bytes()
59    }
60
61    pub fn as_str(&self) -> &str {
62        match self {
63            SampleEntryType::Mp4a => SAMPLE_ENTRY_MP4A.as_str(),
64            SampleEntryType::Aavd => SAMPLE_ENTRY_AAVD.as_str(),
65            SampleEntryType::Text => SAMPLE_ENTRY_TEXT.as_str(),
66            SampleEntryType::Unknown(bytes) => bytes.as_str(),
67        }
68    }
69}
70
71#[derive(Debug, Clone)]
72pub enum SampleEntryData {
73    Audio(AudioSampleEntry),
74    Text(TextSampleEntry),
75    Other(Vec<u8>),
76}
77
78#[derive(Debug, Clone)]
79pub struct SampleEntry {
80    /// Sample entry type (4CC code)
81    pub entry_type: SampleEntryType,
82    /// Data reference index
83    pub data_reference_index: u16,
84    /// Raw sample entry data (codec-specific)
85    pub data: SampleEntryData,
86}
87
88#[derive(Default, Debug, Clone)]
89pub struct SampleDescriptionTableAtom {
90    /// Version of the stsd atom format (0)
91    pub version: u8,
92    /// Flags for the stsd atom (usually all zeros)
93    pub flags: [u8; 3],
94    /// List of sample entries
95    pub entries: Vec<SampleEntry>,
96}
97
98impl From<Vec<SampleEntry>> for SampleDescriptionTableAtom {
99    fn from(entries: Vec<SampleEntry>) -> Self {
100        SampleDescriptionTableAtom {
101            version: 0,
102            flags: [0u8; 3],
103            entries,
104        }
105    }
106}
107
108impl SampleDescriptionTableAtom {
109    pub fn find_or_create_entry<P, D>(&mut self, pred: P, default_fn: D) -> &mut SampleEntry
110    where
111        P: Fn(&SampleEntry) -> bool,
112        D: FnOnce() -> SampleEntry,
113    {
114        if let Some(index) = self.entries.iter().position(pred) {
115            return &mut self.entries[index];
116        }
117        self.entries.push(default_fn());
118        self.entries.last_mut().unwrap()
119    }
120}
121
122impl ParseAtomData for SampleDescriptionTableAtom {
123    fn parse_atom_data(atom_type: FourCC, input: &[u8]) -> Result<Self, ParseError> {
124        crate::atom::util::parser::assert_atom_type!(atom_type, STSD);
125        use crate::atom::util::parser::stream;
126        use winnow::Parser;
127        Ok(parser::parse_stsd_data.parse(stream(input))?)
128    }
129}
130
131impl SerializeAtom for SampleDescriptionTableAtom {
132    fn atom_type(&self) -> FourCC {
133        STSD
134    }
135
136    fn into_body_bytes(self) -> Vec<u8> {
137        serializer::serialize_stsd_atom(self)
138    }
139}
140
141mod serializer {
142    use super::{
143        audio::serializer::serialize_audio_sample_entry,
144        text::serializer::serialize_text_sample_entry, SampleDescriptionTableAtom, SampleEntryData,
145    };
146    use crate::atom::util::serializer::{be_u32, prepend_size_inclusive, SizeU32};
147
148    pub fn serialize_stsd_atom(stsd: SampleDescriptionTableAtom) -> Vec<u8> {
149        let mut data = Vec::new();
150
151        data.push(stsd.version);
152        data.extend(stsd.flags);
153        data.extend(be_u32(
154            stsd.entries
155                .len()
156                .try_into()
157                .expect("stsd entries len must fit in u32"),
158        ));
159
160        for entry in stsd.entries {
161            data.extend(prepend_size_inclusive::<SizeU32, _>(move || {
162                let mut entry_data = Vec::new();
163                entry_data.extend(entry.entry_type.into_bytes());
164                entry_data.extend([0u8; 6]); // reserved
165                entry_data.extend(entry.data_reference_index.to_be_bytes());
166                match entry.data {
167                    SampleEntryData::Audio(audio) => {
168                        entry_data.extend(serialize_audio_sample_entry(audio));
169                    }
170                    SampleEntryData::Text(text) => {
171                        entry_data.extend(serialize_text_sample_entry(text));
172                    }
173                    SampleEntryData::Other(other_data) => {
174                        entry_data.extend_from_slice(&other_data);
175                    }
176                }
177                entry_data
178            }));
179        }
180
181        data
182    }
183}
184
185mod parser {
186    use winnow::{
187        binary::{be_u16, be_u32, length_repeat},
188        combinator::seq,
189        error::StrContext,
190        ModalResult, Parser,
191    };
192
193    use super::{
194        audio::parser::parse_audio_sample_entry, text::parser::parse_text_sample_entry,
195        SampleDescriptionTableAtom, SampleEntry, SampleEntryData, SampleEntryType,
196    };
197    use crate::atom::util::parser::{
198        byte_array, combinators::inclusive_length_and_then, flags3, fourcc, rest_vec, version,
199        Stream,
200    };
201
202    pub fn parse_stsd_data(input: &mut Stream<'_>) -> ModalResult<SampleDescriptionTableAtom> {
203        seq!(SampleDescriptionTableAtom {
204            version: version.verify(|v| *v == 0),
205            flags: flags3,
206            entries: length_repeat(be_u32, parse_sample_entry)
207                .context(StrContext::Label("entries")),
208        })
209        .parse_next(input)
210    }
211
212    fn parse_sample_entry(input: &mut Stream<'_>) -> ModalResult<SampleEntry> {
213        inclusive_length_and_then(
214            be_u32,
215            seq!(SampleEntry {
216                entry_type: fourcc
217                    .map(SampleEntryType::from)
218                    .context(StrContext::Label("entry_type")),
219                _: byte_array::<6>.context(StrContext::Label("reserved")), // reserved
220                data_reference_index: be_u16.context(StrContext::Label("data_reference_index")),
221                data: match entry_type {
222                    SampleEntryType::Mp4a | SampleEntryType::Aavd => {
223                        parse_audio_sample_entry
224                    }
225                    SampleEntryType::Text => {
226                        parse_text_sample_entry
227                    }
228                    _ => parse_unknown_sample_entry,
229                }.context(StrContext::Label("data")),
230            }),
231        )
232        .parse_next(input)
233    }
234
235    pub fn parse_unknown_sample_entry(input: &mut Stream<'_>) -> ModalResult<SampleEntryData> {
236        rest_vec.map(SampleEntryData::Other).parse_next(input)
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use crate::atom::test_utils::test_atom_roundtrip;
243
244    use super::*;
245
246    /// Test round-trip for all available stsd test data files
247    #[test]
248    fn test_stsd_roundtrip() {
249        test_atom_roundtrip::<SampleDescriptionTableAtom>(STSD);
250    }
251}