atmos/os_lib/

fmp4.rs

// fmp4.rs - 断片化 MP4（fMP4、ISO/IEC 14496-12 の Movie Fragment）のデマルチプレクサ。純粋モジュール。
//
// 【2026-10-17】YouTube の DASH / MSE は映像・音声とも fMP4 のセグメント（初期化セグメント =
// ftyp + moov、メディアセグメント = [styp] [sidx] moof + mdat）で届くため新設。`SourceBuffer.appendBuffer`
// のように任意の位置で区切られたバイト列を `push` で受け取り、完結したトップレベルのボックスから順に
// 処理して、時刻付きのサンプル（復号順）を返す。
//
// 対応: moov（trak の tkhd・mdhd・hdlr・stsd の avcC/esds、mvex/trex の既定値）、moof（mfhd・traf の
// tfhd・tfdt・trun。default-base-is-moof・base-data-offset・既定値の継承）、mdat、64 ビット長のボックス。
// 未対応: 暗号化（senc/saiz）、サンプルグループ、edit list による時刻のずらし。

use super::mp4::{extract_asc_from_stsd, extract_avcc_from_stsd};
use alloc::vec::Vec;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TrackKind {
    Video,
    Audio,
    Other,
}

/// 初期化セグメント（moov）から分かるトラックの情報。
#[derive(Clone, Debug)]
pub struct TrackInfo {
    pub track_id: u32,
    pub kind: TrackKind,
    /// mdhd の timescale（1 秒あたりの刻み）。
    pub timescale: u32,
    /// 映像の幅・高さ（ピクセル。音声トラックまたは不明なら 0）。
    pub width: u32,
    pub height: u32,
    /// H.264 の avcC（映像）。
    pub avcc: Option<Vec<u8>>,
    /// AAC の AudioSpecificConfig（音声）。
    pub asc: Option<Vec<u8>>,
    default_duration: u32,
    default_size: u32,
    default_flags: u32,
}

/// 1 サンプル（映像なら 1 アクセスユニット、AAC なら 1 フレーム）。時刻は timescale 単位。
#[derive(Clone, Debug)]
pub struct FragSample {
    pub track_id: u32,
    /// 復号時刻。
    pub dts: u64,
    /// 表示時刻（dts + composition offset）。
    pub pts: i64,
    pub duration: u32,
    pub keyframe: bool,
    pub data: Vec<u8>,
}

/// trun から求めた、mdat の中のサンプルの位置（ストリーム先頭からの絶対位置）。
struct PendingSample {
    track_id: u32,
    dts: u64,
    pts: i64,
    duration: u32,
    keyframe: bool,
    offset: u64,
    size: u32,
}

/// fMP4 デマルチプレクサ。バイト列を順に `push` する。
#[derive(Default)]
pub struct Fmp4Demuxer {
    pub tracks: Vec<TrackInfo>,
    /// まだ完結していないボックスのバイト列（`buf_pos` はその先頭のストリーム上の位置）。
    buf: Vec<u8>,
    buf_pos: u64,
    /// 直前の moof から求めたサンプル（次の mdat で中身を取り出す）。
    pending: Vec<PendingSample>,
    /// トラックごとの次の復号時刻（tfdt が無いフラグメントの続きに使う）。
    next_dts: Vec<(u32, u64)>,
}

fn be32(d: &[u8], p: usize) -> Option<u32> {
    d.get(p..p + 4).map(|b| u32::from_be_bytes([b[0], b[1], b[2], b[3]]))
}

fn be64(d: &[u8], p: usize) -> Option<u64> {
    d.get(p..p + 8).map(|b| u64::from_be_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]))
}

/// ボックスの (ヘッダ長, 全体長, 種類)。全体長 0 は「末尾まで」。
fn box_header(d: &[u8], p: usize) -> Option<(usize, u64, [u8; 4])> {
    let size = be32(d, p)? as u64;
    let t = d.get(p + 4..p + 8)?;
    let typ = [t[0], t[1], t[2], t[3]];
    if size == 1 {
        Some((16, be64(d, p + 8)?, typ))
    } else {
        Some((8, size, typ))
    }
}

/// `d` の子ボックスを (種類, 中身) で列挙する。
fn children(d: &[u8]) -> Vec<([u8; 4], &[u8])> {
    let mut out = Vec::new();
    let mut p = 0;
    while let Some((hdr, size, typ)) = box_header(d, p) {
        let size = if size == 0 { (d.len() - p) as u64 } else { size };
        if size < hdr as u64 || p as u64 + size > d.len() as u64 {
            break;
        }
        out.push((typ, &d[p + hdr..p + size as usize]));
        p += size as usize;
    }
    out
}

fn child<'a>(d: &'a [u8], typ: &[u8; 4]) -> Option<&'a [u8]> {
    children(d).into_iter().find(|(t, _)| t == typ).map(|(_, c)| c)
}

impl Fmp4Demuxer {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn track(&self, id: u32) -> Option<&TrackInfo> {
        self.tracks.iter().find(|t| t.track_id == id)
    }

    /// バイト列を追加し、取り出せたサンプル（復号順）を返す。
    pub fn push(&mut self, bytes: &[u8]) -> Result<Vec<FragSample>, &'static str> {
        self.buf.extend_from_slice(bytes);
        let mut out = Vec::new();
        while let Some((hdr, size, typ)) = box_header(&self.buf, 0) {
            if size == 0 {
                // 「末尾まで」の mdat は区切りが分からないので、来た分を末尾とみなさず待つ。
                break;
            }
            if size < hdr as u64 {
                return Err("fMP4: bad box size");
            }
            if size > 512 * 1024 * 1024 {
                return Err("fMP4: box too large");
            }
            if (self.buf.len() as u64) < size {
                break;
            }
            let size = size as usize;
            let start = self.buf_pos;
            let body: Vec<u8> = self.buf[hdr..size].to_vec();
            match &typ {
                b"moov" => self.parse_moov(&body)?,
                b"moof" => self.parse_moof(&body, start)?,
                b"mdat" => self.take_mdat(&body, start + hdr as u64, &mut out),
                _ => {} // ftyp / styp / sidx / free など
            }
            self.buf.drain(..size);
            self.buf_pos += size as u64;
        }
        Ok(out)
    }

    fn parse_moov(&mut self, moov: &[u8]) -> Result<(), &'static str> {
        self.tracks.clear();
        for (t, trak) in children(moov) {
            if &t != b"trak" {
                continue;
            }
            let Some(tkhd) = child(trak, b"tkhd") else { continue };
            let is_tkhd_v1 = tkhd.first() == Some(&1);
            let track_id = if is_tkhd_v1 { be32(tkhd, 20) } else { be32(tkhd, 12) }.unwrap_or(0);
            // tkhd 末尾の width / height（16.16 固定小数点、v0 は offset 76/80、v1 は 88/92）
            let tkhd_w = be32(tkhd, if is_tkhd_v1 { 88 } else { 76 }).map_or(0, |v| v >> 16);
            let tkhd_h = be32(tkhd, if is_tkhd_v1 { 92 } else { 80 }).map_or(0, |v| v >> 16);
            let Some(mdia) = child(trak, b"mdia") else { continue };
            let timescale = child(mdia, b"mdhd")
                .and_then(|m| if m.first() == Some(&1) { be32(m, 20) } else { be32(m, 12) })
                .unwrap_or(1000)
                .max(1);
            let kind = match child(mdia, b"hdlr").and_then(|h| h.get(8..12)) {
                Some(b"vide") => TrackKind::Video,
                Some(b"soun") => TrackKind::Audio,
                _ => TrackKind::Other,
            };
            let stsd = child(mdia, b"minf").and_then(|m| child(m, b"stbl")).and_then(|s| child(s, b"stsd"));
            // stsd 内の VisualSampleEntry (offset 8+16+16 = 40..44 に u16 width, u16 height)
            let (stsd_w, stsd_h) = stsd
                .and_then(|s| s.get(40..44))
                .map(|b| (u16::from_be_bytes([b[0], b[1]]) as u32, u16::from_be_bytes([b[2], b[3]]) as u32))
                .unwrap_or((0, 0));
            let (width, height) = if kind == TrackKind::Video {
                (if stsd_w > 0 { stsd_w } else { tkhd_w }, if stsd_h > 0 { stsd_h } else { tkhd_h })
            } else {
                (0, 0)
            };
            let (avcc, asc) = match (kind, stsd) {
                (TrackKind::Video, Some(s)) => (extract_avcc_from_stsd(s), None),
                (TrackKind::Audio, Some(s)) => (None, extract_asc_from_stsd(s)),
                _ => (None, None),
            };
            self.tracks.push(TrackInfo {
                track_id,
                kind,
                timescale,
                width,
                height,
                avcc,
                asc,
                default_duration: 0,
                default_size: 0,
                default_flags: 0,
            });
        }
        // mvex/trex（トラックごとの既定値）
        if let Some(mvex) = child(moov, b"mvex") {
            for (t, trex) in children(mvex) {
                if &t != b"trex" {
                    continue;
                }
                let id = be32(trex, 4).unwrap_or(0);
                if let Some(tr) = self.tracks.iter_mut().find(|x| x.track_id == id) {
                    tr.default_duration = be32(trex, 12).unwrap_or(0);
                    tr.default_size = be32(trex, 16).unwrap_or(0);
                    tr.default_flags = be32(trex, 20).unwrap_or(0);
                }
            }
        }
        if self.tracks.is_empty() {
            return Err("fMP4: no tracks in moov");
        }
        Ok(())
    }

    /// moof（`moof_pos` はその先頭のストリーム上の位置）から各サンプルの位置と時刻を求める。
    fn parse_moof(&mut self, moof: &[u8], moof_pos: u64) -> Result<(), &'static str> {
        self.pending.clear();
        // base-data-offset も default-base-is-moof も無いときの既定: 最初の traf は moof の先頭、
        // 以降は直前の traf のデータの終わり（8.8.7.1）。
        let mut implicit_base = moof_pos;
        for (t, traf) in children(moof) {
            if &t != b"traf" {
                continue;
            }
            let tfhd = child(traf, b"tfhd").ok_or("fMP4: traf without tfhd")?;
            let flags = be32(tfhd, 0).unwrap_or(0) & 0xff_ffff;
            let track_id = be32(tfhd, 4).ok_or("fMP4: short tfhd")?;
            let tr = self.track(track_id).cloned();
            let (mut dd, mut ds, mut df) = tr.as_ref().map_or((0, 0, 0), |t| (t.default_duration, t.default_size, t.default_flags));
            let mut p = 8;
            let mut base = if flags & 0x02_0000 != 0 { moof_pos } else { implicit_base };
            if flags & 0x01 != 0 {
                base = be64(tfhd, p).ok_or("fMP4: short tfhd")?;
                p += 8;
            }
            if flags & 0x02 != 0 {
                p += 4; // sample_description_index
            }
            if flags & 0x08 != 0 {
                dd = be32(tfhd, p).ok_or("fMP4: short tfhd")?;
                p += 4;
            }
            if flags & 0x10 != 0 {
                ds = be32(tfhd, p).ok_or("fMP4: short tfhd")?;
                p += 4;
            }
            if flags & 0x20 != 0 {
                df = be32(tfhd, p).ok_or("fMP4: short tfhd")?;
            }
            let mut dts = match child(traf, b"tfdt") {
                Some(tfdt) if tfdt.first() == Some(&1) => be64(tfdt, 4).unwrap_or(0),
                Some(tfdt) => be32(tfdt, 4).unwrap_or(0) as u64,
                None => self.next_dts.iter().find(|(id, _)| *id == track_id).map_or(0, |x| x.1),
            };
            let mut data_end = base;
            for (t2, trun) in children(traf) {
                if &t2 != b"trun" {
                    continue;
                }
                let version = trun.first().copied().unwrap_or(0);
                let tf = be32(trun, 0).unwrap_or(0) & 0xff_ffff;
                let count = be32(trun, 4).ok_or("fMP4: short trun")? as usize;
                if count > 1_000_000 {
                    return Err("fMP4: too many samples in trun");
                }
                let mut q = 8;
                let mut off = data_end;
                if tf & 0x01 != 0 {
                    let d = be32(trun, q).ok_or("fMP4: short trun")? as i32;
                    off = (base as i64 + d as i64).max(0) as u64;
                    q += 4;
                }
                let mut first_flags = None;
                if tf & 0x04 != 0 {
                    first_flags = Some(be32(trun, q).ok_or("fMP4: short trun")?);
                    q += 4;
                }
                for i in 0..count {
                    let mut dur = dd;
                    let mut size = ds;
                    let mut sflags = if i == 0 { first_flags.unwrap_or(df) } else { df };
                    let mut cto: i64 = 0;
                    if tf & 0x100 != 0 {
                        dur = be32(trun, q).ok_or("fMP4: short trun")?;
                        q += 4;
                    }
                    if tf & 0x200 != 0 {
                        size = be32(trun, q).ok_or("fMP4: short trun")?;
                        q += 4;
                    }
                    if tf & 0x400 != 0 {
                        let f = be32(trun, q).ok_or("fMP4: short trun")?;
                        if !(i == 0 && first_flags.is_some()) {
                            sflags = f;
                        }
                        q += 4;
                    }
                    if tf & 0x800 != 0 {
                        let v = be32(trun, q).ok_or("fMP4: short trun")?;
                        cto = if version == 0 { v as i64 } else { v as i32 as i64 };
                        q += 4;
                    }
                    // sample_is_non_sync_sample（ビット 16）が 0 ならキーフレーム。
                    let keyframe = sflags & 0x1_0000 == 0;
                    self.pending.push(PendingSample {
                        track_id,
                        dts,
                        pts: dts as i64 + cto,
                        duration: dur,
                        keyframe,
                        offset: off,
                        size,
                    });
                    dts += dur as u64;
                    off += size as u64;
                }
                data_end = off;
            }
            implicit_base = data_end;
            if let Some(e) = self.next_dts.iter_mut().find(|(id, _)| *id == track_id) {
                e.1 = dts;
            } else {
                self.next_dts.push((track_id, dts));
            }
        }
        Ok(())
    }

    /// mdat（中身の先頭が `data_pos`）から、直前の moof のサンプルを取り出す。
    fn take_mdat(&mut self, mdat: &[u8], data_pos: u64, out: &mut Vec<FragSample>) {
        let end = data_pos + mdat.len() as u64;
        self.pending.retain(|s| {
            if s.offset >= data_pos && s.offset + s.size as u64 <= end {
                let a = (s.offset - data_pos) as usize;
                out.push(FragSample {
                    track_id: s.track_id,
                    dts: s.dts,
                    pts: s.pts,
                    duration: s.duration,
                    keyframe: s.keyframe,
                    data: mdat[a..a + s.size as usize].to_vec(),
                });
                false
            } else {
                true
            }
        });
    }
}