mp4-track 0.1.0

//! Sample-count limits and structural flags on synthesized files.

#![allow(unsafe_code)]

use mp4_track::{Error, Mp4Reader};
use std::alloc::{GlobalAlloc, Layout, System};
use std::io::Cursor;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;

/// The allocation counters are process-wide, so every test in this binary
/// serializes on one lock to keep measurements attributable.
static SERIAL: Mutex<()> = Mutex::new(());

fn serial_guard() -> std::sync::MutexGuard<'static, ()> {
    SERIAL
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
}

struct CountingAllocator;

static CURRENT: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

fn record_allocation(size: usize) {
    let current = CURRENT
        .fetch_add(size, Ordering::Relaxed)
        .saturating_add(size);
    PEAK.fetch_max(current, Ordering::Relaxed);
}

unsafe impl GlobalAlloc for CountingAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let pointer = unsafe { System.alloc(layout) };
        if !pointer.is_null() {
            record_allocation(layout.size());
        }
        pointer
    }

    unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
        CURRENT.fetch_sub(layout.size(), Ordering::Relaxed);
        unsafe { System.dealloc(pointer, layout) };
    }

    unsafe fn realloc(&self, pointer: *mut u8, old: Layout, new_size: usize) -> *mut u8 {
        let resized = unsafe { System.realloc(pointer, old, new_size) };
        if !resized.is_null() {
            if new_size >= old.size() {
                record_allocation(new_size - old.size());
            } else {
                CURRENT.fetch_sub(old.size() - new_size, Ordering::Relaxed);
            }
        }
        resized
    }
}

#[global_allocator]
static ALLOCATOR: CountingAllocator = CountingAllocator;

fn boxed(kind: &[u8; 4], body: &[u8]) -> Vec<u8> {
    let mut result = Vec::with_capacity(body.len() + 8);
    result.extend_from_slice(
        &u32::try_from(body.len() + 8)
            .expect("box size")
            .to_be_bytes(),
    );
    result.extend_from_slice(kind);
    result.extend_from_slice(body);
    result
}

fn full(kind: &[u8; 4], version: u8, flags: u32, body: &[u8]) -> Vec<u8> {
    let mut inner = vec![version];
    inner.extend_from_slice(&flags.to_be_bytes()[1..]);
    inner.extend_from_slice(body);
    boxed(kind, &inner)
}

fn minimal_stbl(stsz: Vec<u8>) -> Vec<u8> {
    let mut entry = Vec::new();
    entry.extend_from_slice(&16u32.to_be_bytes());
    entry.extend_from_slice(b"abcd");
    entry.extend_from_slice(&[0u8; 8]);
    let mut stsd_body = 1u32.to_be_bytes().to_vec();
    stsd_body.extend_from_slice(&entry);
    let stsd = full(b"stsd", 0, 0, &stsd_body);
    let mut stts_body = 1u32.to_be_bytes().to_vec();
    stts_body.extend_from_slice(&1u32.to_be_bytes());
    stts_body.extend_from_slice(&100u32.to_be_bytes());
    let stts = full(b"stts", 0, 0, &stts_body);
    let mut stsc_body = 1u32.to_be_bytes().to_vec();
    for value in [1u32, 1, 1] {
        stsc_body.extend_from_slice(&value.to_be_bytes());
    }
    let stsc = full(b"stsc", 0, 0, &stsc_body);
    let mut stco_body = 1u32.to_be_bytes().to_vec();
    stco_body.extend_from_slice(&0u32.to_be_bytes());
    let stco = full(b"stco", 0, 0, &stco_body);
    let mut body = Vec::new();
    for child in [stsd, stts, stsc, stsz, stco] {
        body.extend_from_slice(&child);
    }
    boxed(b"stbl", &body)
}

fn minimal_moov(stsz: Vec<u8>, with_mvex: bool) -> Vec<u8> {
    let mut mvhd_body = [0u8; 8].to_vec();
    mvhd_body.extend_from_slice(&1000u32.to_be_bytes());
    mvhd_body.extend_from_slice(&0u32.to_be_bytes());
    mvhd_body.extend_from_slice(&0x0001_0000u32.to_be_bytes());
    mvhd_body.extend_from_slice(&0x0100u16.to_be_bytes());
    mvhd_body.extend_from_slice(&[0u8; 10]);
    for value in [0x0001_0000u32, 0, 0, 0, 0x0001_0000, 0, 0, 0, 0x4000_0000] {
        mvhd_body.extend_from_slice(&value.to_be_bytes());
    }
    mvhd_body.extend_from_slice(&[0u8; 24]);
    mvhd_body.extend_from_slice(&2u32.to_be_bytes());
    let mvhd = full(b"mvhd", 0, 0, &mvhd_body);
    let mut tkhd_body = [0u8; 8].to_vec();
    tkhd_body.extend_from_slice(&1u32.to_be_bytes());
    tkhd_body.extend_from_slice(&[0u8; 68]);
    let tkhd = full(b"tkhd", 0, 1, &tkhd_body);
    let mut mdhd_body = [0u8; 8].to_vec();
    mdhd_body.extend_from_slice(&48_000u32.to_be_bytes());
    mdhd_body.extend_from_slice(&0u32.to_be_bytes());
    mdhd_body.extend_from_slice(&0x55c4u16.to_be_bytes());
    mdhd_body.extend_from_slice(&0u16.to_be_bytes());
    let mdhd = full(b"mdhd", 0, 0, &mdhd_body);
    let mut hdlr_body = [0u8; 4].to_vec();
    hdlr_body.extend_from_slice(b"soun");
    hdlr_body.extend_from_slice(&[0u8; 12]);
    hdlr_body.extend_from_slice(b"x\0");
    let hdlr = full(b"hdlr", 0, 0, &hdlr_body);
    let url = full(b"url ", 0, 1, &[]);
    let mut dref_body = 1u32.to_be_bytes().to_vec();
    dref_body.extend_from_slice(&url);
    let dinf = boxed(b"dinf", &full(b"dref", 0, 0, &dref_body));
    let smhd = full(b"smhd", 0, 0, &[0u8; 4]);
    let mut minf_body = Vec::new();
    for child in [smhd, dinf, minimal_stbl(stsz)] {
        minf_body.extend_from_slice(&child);
    }
    let minf = boxed(b"minf", &minf_body);
    let mut mdia_body = Vec::new();
    for child in [mdhd, hdlr, minf] {
        mdia_body.extend_from_slice(&child);
    }
    let mdia = boxed(b"mdia", &mdia_body);
    let mut trak_body = tkhd;
    trak_body.extend_from_slice(&mdia);
    let trak = boxed(b"trak", &trak_body);
    let mut moov_body = mvhd;
    moov_body.extend_from_slice(&trak);
    if with_mvex {
        let mut trex_body = Vec::new();
        for value in [1u32, 1, 0, 0, 0] {
            trex_body.extend_from_slice(&value.to_be_bytes());
        }
        let mvex = boxed(b"mvex", &full(b"trex", 0, 0, &trex_body));
        moov_body.extend_from_slice(&mvex);
    }
    boxed(b"moov", &moov_body)
}

fn minimal_file(stsz: Vec<u8>, with_mvex: bool) -> Vec<u8> {
    let mut ftyp_body = b"isom".to_vec();
    ftyp_body.extend_from_slice(&512u32.to_be_bytes());
    ftyp_body.extend_from_slice(b"isom");
    let mut file = boxed(b"ftyp", &ftyp_body);
    file.extend_from_slice(&minimal_moov(stsz, with_mvex));
    file.extend_from_slice(&boxed(b"mdat", &[0u8; 16]));
    file
}

fn one_sample_stsz() -> Vec<u8> {
    let mut body = 0u32.to_be_bytes().to_vec();
    body.extend_from_slice(&1u32.to_be_bytes());
    body.extend_from_slice(&10u32.to_be_bytes());
    full(b"stsz", 0, 0, &body)
}

fn open_expecting_sample_count_limit(file: Vec<u8>, allocation_bound: usize) {
    let baseline = CURRENT.load(Ordering::Relaxed);
    PEAK.store(baseline, Ordering::Relaxed);
    let result = Mp4Reader::open(Cursor::new(file));
    let peak = PEAK.load(Ordering::Relaxed).saturating_sub(baseline);
    assert!(
        matches!(
            result,
            Err(Error::LimitExceeded {
                limit: "max_sample_count",
                ..
            })
        ),
        "expected the sample-count limit"
    );
    assert!(
        peak <= allocation_bound,
        "peak {peak} exceeds {allocation_bound}"
    );
}

#[test]
fn b_stsz_fixed_count_bomb_rejects_before_allocation() {
    let _serial = serial_guard();
    // stsz with sample_size 1 and a declared count of one million carries no
    // per-sample data, so only the count limit can stop the expansion.
    let mut body = 1u32.to_be_bytes().to_vec();
    body.extend_from_slice(&1_000_000u32.to_be_bytes());
    let stsz = full(b"stsz", 0, 0, &body);
    let file = minimal_file(stsz, false);
    assert!(file.len() < 1024, "the bomb must stay tiny: {}", file.len());
    open_expecting_sample_count_limit(file, 1024 * 1024);
}

#[test]
fn b_stsz_variable_count_bomb_rejects_before_expansion() {
    let _serial = serial_guard();
    // Four million declared sizes fit a 16 MiB table but must not expand
    // into hundreds of megabytes of per-sample records.
    let count = 4_000_000u32;
    let mut body = 0u32.to_be_bytes().to_vec();
    body.extend_from_slice(&count.to_be_bytes());
    body.resize(body.len() + usize::try_from(count).expect("count") * 4, 0);
    let stsz = full(b"stsz", 0, 0, &body);
    let file = minimal_file(stsz, false);
    open_expecting_sample_count_limit(file, 64 * 1024 * 1024);
}

#[test]
fn b_trun_count_without_fields_rejects_before_allocation() {
    let _serial = serial_guard();
    let mut file = minimal_file(one_sample_stsz(), true);
    let mfhd = full(b"mfhd", 0, 0, &1u32.to_be_bytes());
    let tfhd = full(b"tfhd", 0, 0x020000, &1u32.to_be_bytes());
    let trun = full(b"trun", 0, 0, &0xFFFF_0000u32.to_be_bytes());
    let mut traf_body = tfhd;
    traf_body.extend_from_slice(&trun);
    let traf = boxed(b"traf", &traf_body);
    let mut moof_body = mfhd;
    moof_body.extend_from_slice(&traf);
    file.extend_from_slice(&boxed(b"moof", &moof_body));
    open_expecting_sample_count_limit(file, 1024 * 1024);
}

#[test]
fn mvex_without_moof_marks_fragmented() {
    let _serial = serial_guard();
    let file = minimal_file(one_sample_stsz(), true);
    let reader = Mp4Reader::open(Cursor::new(file)).expect("synthetic file must open");
    assert!(reader.fragments().is_empty());
    assert!(reader.is_fragmented());
}

#[test]
fn raised_sample_count_limit_accepts_larger_tables() {
    let _serial = serial_guard();
    let mut body = 1u32.to_be_bytes().to_vec();
    body.extend_from_slice(&1_000_000u32.to_be_bytes());
    let stsz = full(b"stsz", 0, 0, &body);
    let file = minimal_file(stsz, false);
    let mut limits = mp4_track::Limits::default();
    limits.max_sample_count = 2_000_000;
    // The tables stay internally inconsistent on purpose; the point is that
    // the failure moves past the count limit once the caller raises it.
    let result = Mp4Reader::open_with(Cursor::new(file), limits);
    assert!(!matches!(
        result,
        Err(Error::LimitExceeded {
            limit: "max_sample_count",
            ..
        })
    ));
}