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    /// 音声トラックの AudioSpecificConfig(`mp4a` の `esds` の DecoderSpecificInfo)。AAC の初期化に使う。
22    pub audio_asc: Option<Vec<u8>>,
23}
24
25impl Mp4Demuxer {
26    pub fn parse(data: &[u8]) -> Result<Self, &'static str> {
27        let mut demuxer = Self {
28            video_samples: Vec::new(),
29            audio_samples: Vec::new(),
30            duration_ms: 0,
31            video_avcc: None,
32            audio_asc: None,
33        };
34
35        let mut offset = 0;
36        while offset + 8 <= data.len() {
37            let size = u32::from_be_bytes([
38                data[offset],
39                data[offset + 1],
40                data[offset + 2],
41                data[offset + 3],
42            ]) as usize;
43            let box_type = &data[offset + 4..offset + 8];
44
45            if size < 8 {
46                return Err("Invalid box size");
47            }
48
49            if box_type == b"moov" {
50                demuxer.parse_moov(&data[offset + 8..offset + size])?;
51            }
52
53            offset += size;
54        }
55
56        Ok(demuxer)
57    }
58
59    fn parse_moov(&mut self, data: &[u8]) -> Result<(), &'static str> {
60        let mut offset = 0;
61        while offset + 8 <= data.len() {
62            let size = u32::from_be_bytes([
63                data[offset],
64                data[offset + 1],
65                data[offset + 2],
66                data[offset + 3],
67            ]) as usize;
68            let box_type = &data[offset + 4..offset + 8];
69
70            if size < 8 || offset + size > data.len() {
71                return Err("Invalid box size in moov");
72            }
73
74            if box_type == b"mvhd" {
75                self.parse_mvhd(&data[offset + 8..offset + size])?;
76            } else if box_type == b"trak" {
77                self.parse_trak(&data[offset + 8..offset + size])?;
78            }
79
80            offset += size;
81        }
82        Ok(())
83    }
84
85    fn parse_mvhd(&mut self, payload: &[u8]) -> Result<(), &'static str> {
86        if payload.is_empty() {
87            return Ok(());
88        }
89        let version = payload[0];
90        if version == 0 {
91            if payload.len() >= 20 {
92                let timescale = u32::from_be_bytes([payload[12], payload[13], payload[14], payload[15]]);
93                let duration = u32::from_be_bytes([payload[16], payload[17], payload[18], payload[19]]);
94                if timescale > 0 {
95                    self.duration_ms = (duration as u64 * 1000) / (timescale as u64);
96                }
97            }
98        } else if version == 1 {
99            if payload.len() >= 32 {
100                let timescale = u32::from_be_bytes([payload[20], payload[21], payload[22], payload[23]]);
101                let duration = u64::from_be_bytes([
102                    payload[24], payload[25], payload[26], payload[27],
103                    payload[28], payload[29], payload[30], payload[31],
104                ]);
105                if timescale > 0 {
106                    self.duration_ms = (duration * 1000) / (timescale as u64);
107                }
108            }
109        }
110        Ok(())
111    }
112
113    fn parse_trak(&mut self, data: &[u8]) -> Result<(), &'static str> {
114        let mut offset = 0;
115        let mut is_video = false;
116        let mut is_audio = false;
117        let mut mdia_data: Option<&[u8]> = None;
118
119        while offset + 8 <= data.len() {
120            let size = u32::from_be_bytes([
121                data[offset],
122                data[offset + 1],
123                data[offset + 2],
124                data[offset + 3],
125            ]) as usize;
126            let box_type = &data[offset + 4..offset + 8];
127
128            if size < 8 || offset + size > data.len() {
129                return Err("Invalid box size in trak");
130            }
131
132            if box_type == b"mdia" {
133                mdia_data = Some(&data[offset + 8..offset + size]);
134            }
135
136            offset += size;
137        }
138
139        if let Some(mdia) = mdia_data {
140            let mut mdia_offset = 0;
141            let mut minf_data: Option<&[u8]> = None;
142            // mdhd の timescale(stts の単位、1 秒あたりの刻み)。無ければ 1000(ミリ秒)とみなす。
143            let mut timescale = 1000u32;
144
145            while mdia_offset + 8 <= mdia.len() {
146                let size = u32::from_be_bytes([
147                    mdia[mdia_offset],
148                    mdia[mdia_offset + 1],
149                    mdia[mdia_offset + 2],
150                    mdia[mdia_offset + 3],
151                ]) as usize;
152                let box_type = &mdia[mdia_offset + 4..mdia_offset + 8];
153
154                if size < 8 || mdia_offset + size > mdia.len() {
155                    return Err("Invalid box size in mdia");
156                }
157
158                if box_type == b"mdhd" {
159                    let p = &mdia[mdia_offset + 8..mdia_offset + size];
160                    let at = if p.first() == Some(&1) { 20 } else { 12 };
161                    if let Some(t) = p.get(at..at + 4) {
162                        timescale = u32::from_be_bytes([t[0], t[1], t[2], t[3]]).max(1);
163                    }
164                } else if box_type == b"hdlr" {
165                    if mdia_offset + 16 <= mdia.len() {
166                        let h_type = &mdia[mdia_offset + 16..mdia_offset + 20];
167                        if h_type == b"vide" {
168                            is_video = true;
169                        } else if h_type == b"soun" {
170                            is_audio = true;
171                        }
172                    }
173                } else if box_type == b"minf" {
174                    minf_data = Some(&mdia[mdia_offset + 8..mdia_offset + size]);
175                }
176
177                mdia_offset += size;
178            }
179
180            if let Some(minf) = minf_data {
181                let mut minf_offset = 0;
182                let mut stbl_data: Option<&[u8]> = None;
183
184                while minf_offset + 8 <= minf.len() {
185                    let size = u32::from_be_bytes([
186                        minf[minf_offset],
187                        minf[minf_offset + 1],
188                        minf[minf_offset + 2],
189                        minf[minf_offset + 3],
190                    ]) as usize;
191                    let box_type = &minf[minf_offset + 4..minf_offset + 8];
192
193                    if size < 8 || minf_offset + size > minf.len() {
194                        return Err("Invalid box size in minf");
195                    }
196
197                    if box_type == b"stbl" {
198                        stbl_data = Some(&minf[minf_offset + 8..minf_offset + size]);
199                    }
200
201                    minf_offset += size;
202                }
203
204                if let Some(stbl) = stbl_data {
205                    if is_video {
206                        self.parse_stbl(stbl, true, timescale)?;
207                    } else if is_audio {
208                        self.parse_stbl(stbl, false, timescale)?;
209                    }
210                }
211            }
212        }
213
214        Ok(())
215    }
216
217    fn parse_stbl(&mut self, data: &[u8], is_video: bool, timescale: u32) -> Result<(), &'static str> {
218        let mut offset = 0;
219        let mut stsz_sizes: Vec<u32> = Vec::new();
220        let mut stco_offsets: Vec<u64> = Vec::new();
221        // stts の (sample_count, sample_delta)。各サンプルの復号時刻(timescale 単位)を作る。
222        let mut stts: Vec<(u32, u32)> = Vec::new();
223        let mut stsc_entries: Vec<(u32, u32, u32)> = Vec::new(); // (first_chunk, samples_per_chunk, sample_description_index)
224        let mut stss_keyframes: Option<Vec<u32>> = None;
225        let mut ctts_entries: Vec<(u32, i32)> = Vec::new();
226
227        while offset + 8 <= data.len() {
228            let size = u32::from_be_bytes([
229                data[offset],
230                data[offset + 1],
231                data[offset + 2],
232                data[offset + 3],
233            ]) as usize;
234            let box_type = &data[offset + 4..offset + 8];
235
236            if size < 8 || offset + size > data.len() {
237                return Err("Invalid box size in stbl");
238            }
239
240            let payload = &data[offset + 8..offset + size];
241
242            if box_type == b"stsd" && is_video {
243                if let Some(avcc) = extract_avcc_from_stsd(payload) {
244                    self.video_avcc = Some(avcc);
245                }
246            } else if box_type == b"stsd" {
247                if let Some(asc) = extract_asc_from_stsd(payload) {
248                    self.audio_asc = Some(asc);
249                }
250            } else if box_type == b"stts" && payload.len() >= 8 {
251                let n = u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]]) as usize;
252                for e in payload[8..].chunks_exact(8).take(n) {
253                    stts.push((
254                        u32::from_be_bytes([e[0], e[1], e[2], e[3]]),
255                        u32::from_be_bytes([e[4], e[5], e[6], e[7]]),
256                    ));
257                }
258            } else if box_type == b"stss" && payload.len() >= 8 {
259                let count = u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]]) as usize;
260                if payload.len() >= 8 + count * 4 {
261                    let mut kf = Vec::with_capacity(count);
262                    for i in 0..count {
263                        let sample_num = u32::from_be_bytes([
264                            payload[8 + i * 4],
265                            payload[8 + i * 4 + 1],
266                            payload[8 + i * 4 + 2],
267                            payload[8 + i * 4 + 3],
268                        ]);
269                        kf.push(sample_num);
270                    }
271                    stss_keyframes = Some(kf);
272                }
273            } else if box_type == b"ctts" && payload.len() >= 8 {
274                let version = payload[0];
275                let count = u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]]) as usize;
276                if payload.len() >= 8 + count * 8 {
277                    ctts_entries.reserve(count);
278                    for i in 0..count {
279                        let s_count = u32::from_be_bytes([
280                            payload[8 + i * 8],
281                            payload[8 + i * 8 + 1],
282                            payload[8 + i * 8 + 2],
283                            payload[8 + i * 8 + 3],
284                        ]);
285                        let s_offset = if version == 0 {
286                            u32::from_be_bytes([
287                                payload[8 + i * 8 + 4],
288                                payload[8 + i * 8 + 5],
289                                payload[8 + i * 8 + 6],
290                                payload[8 + i * 8 + 7],
291                            ]) as i32
292                        } else {
293                            i32::from_be_bytes([
294                                payload[8 + i * 8 + 4],
295                                payload[8 + i * 8 + 5],
296                                payload[8 + i * 8 + 6],
297                                payload[8 + i * 8 + 7],
298                            ])
299                        };
300                        ctts_entries.push((s_count, s_offset));
301                    }
302                }
303            } else if box_type == b"co64" && payload.len() >= 8 {
304                // 64 ビットのチャンクオフセット(4GB を超えるファイル)。
305                let n = u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]]) as usize;
306                for e in payload[8..].chunks_exact(8).take(n) {
307                    stco_offsets.push(u64::from_be_bytes([e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7]]));
308                }
309            } else if box_type == b"stsz" {
310                if payload.len() >= 12 {
311                    let sample_size =
312                        u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]]);
313                    let sample_count =
314                        u32::from_be_bytes([payload[8], payload[9], payload[10], payload[11]])
315                            as usize;
316                    if sample_size > 0 {
317                        stsz_sizes = alloc::vec![sample_size; sample_count];
318                    } else if payload.len() >= 12 + sample_count * 4 {
319                        stsz_sizes.reserve(sample_count);
320                        for i in 0..sample_count {
321                            let s = u32::from_be_bytes([
322                                payload[12 + i * 4],
323                                payload[12 + i * 4 + 1],
324                                payload[12 + i * 4 + 2],
325                                payload[12 + i * 4 + 3],
326                            ]);
327                            stsz_sizes.push(s);
328                        }
329                    }
330                }
331            } else if box_type == b"stco" {
332                if payload.len() >= 8 {
333                    let entry_count =
334                        u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]])
335                            as usize;
336                    if payload.len() >= 8 + entry_count * 4 {
337                        stco_offsets.reserve(entry_count);
338                        for i in 0..entry_count {
339                            let off = u32::from_be_bytes([
340                                payload[8 + i * 4],
341                                payload[8 + i * 4 + 1],
342                                payload[8 + i * 4 + 2],
343                                payload[8 + i * 4 + 3],
344                            ]);
345                            stco_offsets.push(off as u64);
346                        }
347                    }
348                }
349            } else if box_type == b"stsc" && payload.len() >= 8 {
350                let entry_count =
351                    u32::from_be_bytes([payload[4], payload[5], payload[6], payload[7]]) as usize;
352                if payload.len() >= 8 + entry_count * 12 {
353                    for i in 0..entry_count {
354                        let first_chunk = u32::from_be_bytes([
355                            payload[8 + i * 12],
356                            payload[8 + i * 12 + 1],
357                            payload[8 + i * 12 + 2],
358                            payload[8 + i * 12 + 3],
359                        ]);
360                        let samples = u32::from_be_bytes([
361                            payload[8 + i * 12 + 4],
362                            payload[8 + i * 12 + 5],
363                            payload[8 + i * 12 + 6],
364                            payload[8 + i * 12 + 7],
365                        ]);
366                        let desc = u32::from_be_bytes([
367                            payload[8 + i * 12 + 8],
368                            payload[8 + i * 12 + 9],
369                            payload[8 + i * 12 + 10],
370                            payload[8 + i * 12 + 11],
371                        ]);
372                        stsc_entries.push((first_chunk, samples, desc));
373                    }
374                }
375            }
376
377            offset += size;
378        }
379
380        // Chunk offsets and sample sizes を組み合わせて各サンプルの絶対ファイルオフセットを構築する
381        if !stsz_sizes.is_empty() && !stco_offsets.is_empty() && !stsc_entries.is_empty() {
382            // 復号時刻(DTS、timescale 単位)。stts から構築する。
383            let mut dts_times: Vec<u64> = Vec::with_capacity(stsz_sizes.len());
384            let mut t = 0u64;
385            for &(count, delta) in &stts {
386                for _ in 0..count {
387                    if dts_times.len() >= stsz_sizes.len() {
388                        break;
389                    }
390                    dts_times.push(t);
391                    t += delta as u64;
392                }
393            }
394
395            // 表示時刻オフセット(CTTS)。ctts があれば展開。
396            let mut ctts_offsets: Vec<i32> = Vec::with_capacity(stsz_sizes.len());
397            for &(count, offset) in &ctts_entries {
398                for _ in 0..count {
399                    if ctts_offsets.len() >= stsz_sizes.len() {
400                        break;
401                    }
402                    ctts_offsets.push(offset);
403                }
404            }
405
406            let mut sample_idx = 0;
407            let mut current_stsc_idx = 0;
408
409            for chunk_idx in 0..stco_offsets.len() {
410                let chunk_num = (chunk_idx + 1) as u32; // 1-indexed
411
412                // 該当する stsc エントリを探す
413                while current_stsc_idx + 1 < stsc_entries.len()
414                    && chunk_num >= stsc_entries[current_stsc_idx + 1].0
415                {
416                    current_stsc_idx += 1;
417                }
418
419                let samples_in_chunk = stsc_entries[current_stsc_idx].1 as usize;
420                let mut chunk_offset = stco_offsets[chunk_idx] as usize;
421
422                for _ in 0..samples_in_chunk {
423                    if sample_idx >= stsz_sizes.len() {
424                        break;
425                    }
426                    let s_size = stsz_sizes[sample_idx] as usize;
427
428                    let dts = dts_times.get(sample_idx).copied();
429                    let ctts_off = ctts_offsets.get(sample_idx).copied().unwrap_or(0);
430                    let timestamp_ms = match dts {
431                        Some(d) => {
432                            let pts = (d as i64 + ctts_off as i64).max(0) as u64;
433                            pts * 1000 / timescale as u64
434                        }
435                        None => (sample_idx as u64 * 33) + (ctts_off.max(0) as u64 * 1000 / timescale as u64),
436                    };
437
438                    let is_keyframe = match &stss_keyframes {
439                        Some(kf) => {
440                            let sample_num = (sample_idx + 1) as u32; // 1-indexed
441                            kf.binary_search(&sample_num).is_ok()
442                        }
443                        None => true, // stss ボックスが無いトラック(音声等)は全サンプルがキーフレーム
444                    };
445
446                    let sample = Mp4Sample {
447                        offset: chunk_offset,
448                        size: s_size,
449                        is_keyframe,
450                        timestamp_ms,
451                    };
452
453                    if is_video {
454                        self.video_samples.push(sample);
455                    } else {
456                        self.audio_samples.push(sample);
457                    }
458
459                    chunk_offset += s_size;
460                    sample_idx += 1;
461                }
462            }
463        }
464
465        Ok(())
466    }
467}
468
469/// `stsd` ボックスのペイロードから `avc1`/`avc3` エントリ内の `avcC` を取り出す。
470/// stsd payload: [version+flags(4)][entry_count(4)][sample entry...]
471/// sample entry: [size(4)][type(4)][VisualSampleEntry(78)][child boxes...]
472pub(crate) fn extract_avcc_from_stsd(payload: &[u8]) -> Option<Vec<u8>> {
473    if payload.len() < 16 {
474        return None;
475    }
476    let entry_size =
477        u32::from_be_bytes([payload[8], payload[9], payload[10], payload[11]]) as usize;
478    let entry_type = &payload[12..16];
479    if !(entry_type == b"avc1" || entry_type == b"avc3") {
480        return None;
481    }
482    let entry_end = (8 + entry_size).min(payload.len());
483    // child boxes は エントリ box ヘッダ(8) + VisualSampleEntry(78) の後ろから。
484    let mut p = 8 + 8 + 78;
485    while p + 8 <= entry_end {
486        let bsize = u32::from_be_bytes([payload[p], payload[p + 1], payload[p + 2], payload[p + 3]])
487            as usize;
488        let btype = &payload[p + 4..p + 8];
489        if bsize < 8 || p + bsize > entry_end {
490            break;
491        }
492        if btype == b"avcC" {
493            return Some(payload[p + 8..p + bsize].to_vec());
494        }
495        p += bsize;
496    }
497    None
498}
499
500/// `stsd` ボックスのペイロードから `mp4a` エントリの `esds` を探し、AudioSpecificConfig を取り出す。
501/// sample entry: [size(4)][type(4)][予約(6)][data_reference_index(2)][AudioSampleEntry(20)][child boxes...]
502/// QuickTime の版 1・2 の音声エントリはそれぞれ 16・36 バイト長い。
503pub(crate) fn extract_asc_from_stsd(payload: &[u8]) -> Option<Vec<u8>> {
504    if payload.len() < 16 {
505        return None;
506    }
507    let entry_size = u32::from_be_bytes([payload[8], payload[9], payload[10], payload[11]]) as usize;
508    if &payload[12..16] != b"mp4a" {
509        return None;
510    }
511    let entry_end = (8 + entry_size).min(payload.len());
512    let body = 16;
513    let version = u16::from_be_bytes([*payload.get(body + 8)?, *payload.get(body + 9)?]);
514    let extra = match version {
515        1 => 16,
516        2 => 36,
517        _ => 0,
518    };
519    let mut p = body + 28 + extra;
520    while p + 8 <= entry_end {
521        let bsize = u32::from_be_bytes([payload[p], payload[p + 1], payload[p + 2], payload[p + 3]]) as usize;
522        if bsize < 8 || p + bsize > entry_end {
523            break;
524        }
525        if &payload[p + 4..p + 8] == b"esds" {
526            return parse_esds_asc(&payload[p + 8..p + bsize]);
527        }
528        p += bsize;
529    }
530    None
531}
532
533/// `esds`(版・フラグ 4 バイト + ES_Descriptor)から DecoderSpecificInfo(タグ 5)の中身を返す。
534pub fn parse_esds_asc(esds: &[u8]) -> Option<Vec<u8>> {
535    // 記述子の長さ(7 ビットずつ、最上位ビットが継続)。
536    fn desc(d: &[u8], p: &mut usize) -> Option<(u8, usize)> {
537        let tag = *d.get(*p)?;
538        *p += 1;
539        let mut len = 0usize;
540        for _ in 0..4 {
541            let b = *d.get(*p)?;
542            *p += 1;
543            len = (len << 7) | (b & 0x7f) as usize;
544            if b & 0x80 == 0 {
545                break;
546            }
547        }
548        Some((tag, len))
549    }
550    let mut p = 4;
551    let (tag, _) = desc(esds, &mut p)?;
552    if tag != 3 {
553        return None;
554    }
555    p += 2; // ES_ID
556    let flags = *esds.get(p)?;
557    p += 1;
558    if flags & 0x80 != 0 {
559        p += 2; // dependsOn_ES_ID
560    }
561    if flags & 0x40 != 0 {
562        p += 1 + *esds.get(p)? as usize; // URL
563    }
564    if flags & 0x20 != 0 {
565        p += 2; // OCR_ES_Id
566    }
567    let (tag, _) = desc(esds, &mut p)?;
568    if tag != 4 {
569        return None;
570    }
571    p += 13; // objectTypeIndication(1) streamType(1) bufferSize(3) maxBitrate(4) avgBitrate(4)
572    let (tag, len) = desc(esds, &mut p)?;
573    if tag != 5 {
574        return None;
575    }
576    esds.get(p..p + len).map(|v| v.to_vec())
577}