Skip to main content

atmos/os_lib/
mp4.rs

1// mp4.rs - Lightweight no_std MP4 (MPEG-4 Part 14) Demuxer for AtmOS
2#![allow(dead_code)]
3
4use alloc::vec::Vec;
5
6#[derive(Debug, Clone, Copy)]
7pub struct Mp4Sample {
8    pub offset: usize,
9    pub size: usize,
10    pub is_keyframe: bool,
11    pub timestamp_ms: u64,
12}
13
14pub struct Mp4Demuxer {
15    pub video_samples: Vec<Mp4Sample>,
16    pub audio_samples: Vec<Mp4Sample>,
17    pub duration_ms: u64,
18    /// 映像トラックの avcC(AVCDecoderConfigurationRecord)。SPS/PPS と
19    /// NAL 長さフィールドのサイズを内包する。H.264 デコーダの初期化に使う。
20    pub video_avcc: Option<Vec<u8>>,
21}
22
23impl Mp4Demuxer {
24    pub fn parse(data: &[u8]) -> Result<Self, &'static str> {
25        let mut demuxer = Self {
26            video_samples: Vec::new(),
27            audio_samples: Vec::new(),
28            duration_ms: 0,
29            video_avcc: None,
30        };
31
32        let mut offset = 0;
33        while offset + 8 <= data.len() {
34            let size = u32::from_be_bytes([
35                data[offset],
36                data[offset + 1],
37                data[offset + 2],
38                data[offset + 3],
39            ]) as usize;
40            let box_type = &data[offset + 4..offset + 8];
41
42            if size < 8 {
43                return Err("Invalid box size");
44            }
45
46            if box_type == b"moov" {
47                demuxer.parse_moov(&data[offset + 8..offset + size])?;
48            }
49
50            offset += size;
51        }
52
53        Ok(demuxer)
54    }
55
56    fn parse_moov(&mut self, data: &[u8]) -> Result<(), &'static str> {
57        let mut offset = 0;
58        while offset + 8 <= data.len() {
59            let size = u32::from_be_bytes([
60                data[offset],
61                data[offset + 1],
62                data[offset + 2],
63                data[offset + 3],
64            ]) as usize;
65            let box_type = &data[offset + 4..offset + 8];
66
67            if size < 8 || offset + size > data.len() {
68                return Err("Invalid box size in moov");
69            }
70
71            if box_type == b"trak" {
72                self.parse_trak(&data[offset + 8..offset + size])?;
73            }
74
75            offset += size;
76        }
77        Ok(())
78    }
79
80    fn parse_trak(&mut self, data: &[u8]) -> Result<(), &'static str> {
81        let mut offset = 0;
82        let mut is_video = false;
83        let mut is_audio = false;
84        let mut mdia_data: Option<&[u8]> = None;
85
86        while offset + 8 <= data.len() {
87            let size = u32::from_be_bytes([
88                data[offset],
89                data[offset + 1],
90                data[offset + 2],
91                data[offset + 3],
92            ]) as usize;
93            let box_type = &data[offset + 4..offset + 8];
94
95            if size < 8 || offset + size > data.len() {
96                return Err("Invalid box size in trak");
97            }
98
99            if box_type == b"mdia" {
100                mdia_data = Some(&data[offset + 8..offset + size]);
101            }
102
103            offset += size;
104        }
105
106        if let Some(mdia) = mdia_data {
107            let mut mdia_offset = 0;
108            let mut minf_data: Option<&[u8]> = None;
109
110            while mdia_offset + 8 <= mdia.len() {
111                let size = u32::from_be_bytes([
112                    mdia[mdia_offset],
113                    mdia[mdia_offset + 1],
114                    mdia[mdia_offset + 2],
115                    mdia[mdia_offset + 3],
116                ]) as usize;
117                let box_type = &mdia[mdia_offset + 4..mdia_offset + 8];
118
119                if size < 8 || mdia_offset + size > mdia.len() {
120                    return Err("Invalid box size in mdia");
121                }
122
123                if box_type == b"hdlr" {
124                    if mdia_offset + 16 <= mdia.len() {
125                        let h_type = &mdia[mdia_offset + 16..mdia_offset + 20];
126                        if h_type == b"vide" {
127                            is_video = true;
128                        } else if h_type == b"soun" {
129                            is_audio = true;
130                        }
131                    }
132                } else if box_type == b"minf" {
133                    minf_data = Some(&mdia[mdia_offset + 8..mdia_offset + size]);
134                }
135
136                mdia_offset += size;
137            }
138
139            if let Some(minf) = minf_data {
140                let mut minf_offset = 0;
141                let mut stbl_data: Option<&[u8]> = None;
142
143                while minf_offset + 8 <= minf.len() {
144                    let size = u32::from_be_bytes([
145                        minf[minf_offset],
146                        minf[minf_offset + 1],
147                        minf[minf_offset + 2],
148                        minf[minf_offset + 3],
149                    ]) as usize;
150                    let box_type = &minf[minf_offset + 4..minf_offset + 8];
151
152                    if size < 8 || minf_offset + size > minf.len() {
153                        return Err("Invalid box size in minf");
154                    }
155
156                    if box_type == b"stbl" {
157                        stbl_data = Some(&minf[minf_offset + 8..minf_offset + size]);
158                    }
159
160                    minf_offset += size;
161                }
162
163                if let Some(stbl) = stbl_data {
164                    if is_video {
165                        self.parse_stbl(stbl, true)?;
166                    } else if is_audio {
167                        self.parse_stbl(stbl, false)?;
168                    }
169                }
170            }
171        }
172
173        Ok(())
174    }
175
176    fn parse_stbl(&mut self, data: &[u8], is_video: bool) -> Result<(), &'static str> {
177        let mut offset = 0;
178        let mut stsz_sizes: Vec<u32> = Vec::new();
179        let mut stco_offsets: Vec<u32> = Vec::new();
180        let mut stsc_entries: Vec<(u32, u32, u32)> = Vec::new(); // (first_chunk, samples_per_chunk, sample_description_index)
181
182        while offset + 8 <= data.len() {
183            let size = u32::from_be_bytes([
184                data[offset],
185                data[offset + 1],
186                data[offset + 2],
187                data[offset + 3],
188            ]) as usize;
189            let box_type = &data[offset + 4..offset + 8];
190
191            if size < 8 || offset + size > data.len() {
192                return Err("Invalid box size in stbl");
193            }
194
195            let payload = &data[offset + 8..offset + size];
196
197            if box_type == b"stsd" && is_video {
198                if let Some(avcc) = extract_avcc_from_stsd(payload) {
199                    self.video_avcc = Some(avcc);
200                }
201            } else if box_type == b"stsz" {
202                if payload.len() >= 12 {
203                    let sample_size =
204                        u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]]);
205                    let sample_count =
206                        u32::from_be_bytes([payload[8], payload[9], payload[10], payload[11]])
207                            as usize;
208                    if sample_size > 0 {
209                        stsz_sizes = alloc::vec![sample_size; sample_count];
210                    } else if payload.len() >= 12 + sample_count * 4 {
211                        stsz_sizes.reserve(sample_count);
212                        for i in 0..sample_count {
213                            let s = u32::from_be_bytes([
214                                payload[12 + i * 4],
215                                payload[12 + i * 4 + 1],
216                                payload[12 + i * 4 + 2],
217                                payload[12 + i * 4 + 3],
218                            ]);
219                            stsz_sizes.push(s);
220                        }
221                    }
222                }
223            } else if box_type == b"stco" {
224                if payload.len() >= 8 {
225                    let entry_count =
226                        u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]])
227                            as usize;
228                    if payload.len() >= 8 + entry_count * 4 {
229                        stco_offsets.reserve(entry_count);
230                        for i in 0..entry_count {
231                            let off = u32::from_be_bytes([
232                                payload[8 + i * 4],
233                                payload[8 + i * 4 + 1],
234                                payload[8 + i * 4 + 2],
235                                payload[8 + i * 4 + 3],
236                            ]);
237                            stco_offsets.push(off);
238                        }
239                    }
240                }
241            } else if box_type == b"stsc" && payload.len() >= 8 {
242                let entry_count =
243                    u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]]) as usize;
244                if payload.len() >= 8 + entry_count * 12 {
245                    for i in 0..entry_count {
246                        let first_chunk = u32::from_be_bytes([
247                            payload[8 + i * 12],
248                            payload[8 + i * 12 + 1],
249                            payload[8 + i * 12 + 2],
250                            payload[8 + i * 12 + 3],
251                        ]);
252                        let samples = u32::from_be_bytes([
253                            payload[8 + i * 12 + 4],
254                            payload[8 + i * 12 + 5],
255                            payload[8 + i * 12 + 6],
256                            payload[8 + i * 12 + 7],
257                        ]);
258                        let desc = u32::from_be_bytes([
259                            payload[8 + i * 12 + 8],
260                            payload[8 + i * 12 + 9],
261                            payload[8 + i * 12 + 10],
262                            payload[8 + i * 12 + 11],
263                        ]);
264                        stsc_entries.push((first_chunk, samples, desc));
265                    }
266                }
267            }
268
269            offset += size;
270        }
271
272        // Chunk offsets and sample sizes を組み合わせて各サンプルの絶対ファイルオフセットを構築する
273        if !stsz_sizes.is_empty() && !stco_offsets.is_empty() && !stsc_entries.is_empty() {
274            let mut sample_idx = 0;
275            let mut current_stsc_idx = 0;
276
277            for chunk_idx in 0..stco_offsets.len() {
278                let chunk_num = (chunk_idx + 1) as u32; // 1-indexed
279
280                // 該当する stsc エントリを探す
281                while current_stsc_idx + 1 < stsc_entries.len()
282                    && chunk_num >= stsc_entries[current_stsc_idx + 1].0
283                {
284                    current_stsc_idx += 1;
285                }
286
287                let samples_in_chunk = stsc_entries[current_stsc_idx].1 as usize;
288                let mut chunk_offset = stco_offsets[chunk_idx] as usize;
289
290                for _ in 0..samples_in_chunk {
291                    if sample_idx >= stsz_sizes.len() {
292                        break;
293                    }
294                    let s_size = stsz_sizes[sample_idx] as usize;
295
296                    let sample = Mp4Sample {
297                        offset: chunk_offset,
298                        size: s_size,
299                        is_keyframe: true, // 簡易的判断(すべてキーフレーム、または後続のstssで判定)
300                        timestamp_ms: (sample_idx as u64) * 33, // 簡易的に30fps (33ms) または 22.05kHz (1000/22050 ms) と仮定
301                    };
302
303                    if is_video {
304                        self.video_samples.push(sample);
305                    } else {
306                        self.audio_samples.push(sample);
307                    }
308
309                    chunk_offset += s_size;
310                    sample_idx += 1;
311                }
312            }
313        }
314
315        Ok(())
316    }
317}
318
319/// `stsd` ボックスのペイロードから `avc1`/`avc3` エントリ内の `avcC` を取り出す。
320/// stsd payload: [version+flags(4)][entry_count(4)][sample entry...]
321/// sample entry: [size(4)][type(4)][VisualSampleEntry(78)][child boxes...]
322fn extract_avcc_from_stsd(payload: &[u8]) -> Option<Vec<u8>> {
323    if payload.len() < 16 {
324        return None;
325    }
326    let entry_size =
327        u32::from_be_bytes([payload[8], payload[9], payload[10], payload[11]]) as usize;
328    let entry_type = &payload[12..16];
329    if !(entry_type == b"avc1" || entry_type == b"avc3") {
330        return None;
331    }
332    let entry_end = (8 + entry_size).min(payload.len());
333    // child boxes は エントリ box ヘッダ(8) + VisualSampleEntry(78) の後ろから。
334    let mut p = 8 + 8 + 78;
335    while p + 8 <= entry_end {
336        let bsize = u32::from_be_bytes([payload[p], payload[p + 1], payload[p + 2], payload[p + 3]])
337            as usize;
338        let btype = &payload[p + 4..p + 8];
339        if bsize < 8 || p + bsize > entry_end {
340            break;
341        }
342        if btype == b"avcC" {
343            return Some(payload[p + 8..p + bsize].to_vec());
344        }
345        p += bsize;
346    }
347    None
348}