mp4forge 0.8.0

Rust library and CLI for inspecting, probing, extracting, muxing, and rewriting MP4 structures
Documentation
use std::fs::File;
use std::io::Cursor;

#[cfg(feature = "async")]
use tokio::fs::File as TokioFile;

use crate::FourCc;
use crate::bitio::BitReader;
use crate::boxes::avs3::Av3c;

use super::super::MuxError;
use super::super::import::{
    SegmentedMuxSourceSegment, StagedSample, build_btrt_from_sample_sizes,
    build_visual_sample_entry_box,
};
use super::annexb_common::{
    read_bit_labeled, read_bits_u8_labeled, read_bits_u16_labeled, skip_bits_labeled,
};
#[cfg(feature = "async")]
use super::container_common::read_segmented_bytes_async;
use super::container_common::read_segmented_bytes_sync;

const SAMPLE_ENTRY_AVS3: FourCc = FourCc::from_bytes(*b"avs3");
const AVS3_DESCRIPTOR_CONFIG_SIZE: usize = 10;
const AVS3_SEQUENCE_HEADER_PREFIX: [u8; 4] = [0x00, 0x00, 0x01, 0xB0];
const AVS3_INTRA_PICTURE_PREFIX: [u8; 4] = [0x00, 0x00, 0x01, 0xB3];
const AVS3_INTER_PICTURE_PREFIX: [u8; 4] = [0x00, 0x00, 0x01, 0xB6];
const AVS3_PREFIX_SCAN_BYTES: usize = 128;

pub(in crate::mux) struct ParsedTransportAvs3Track {
    pub(in crate::mux) timescale: u32,
    pub(in crate::mux) sample_entry_box: Vec<u8>,
    pub(in crate::mux) samples: Vec<StagedSample>,
}

struct ParsedAvs3SequenceHeader {
    timescale: u32,
    sample_duration: u32,
    low_delay: bool,
}

pub(in crate::mux) fn scan_transport_avs3_segmented_sync(
    file: &mut File,
    segments: &[SegmentedMuxSourceSegment],
    total_size: u64,
    sample_offsets: &[u64],
    avs3_config: &[u8],
    spec: &str,
) -> Result<ParsedTransportAvs3Track, MuxError> {
    let config = parse_transport_avs3_config_bytes(avs3_config, spec)?;
    let sample_bounds = build_transport_avs3_sample_bounds(sample_offsets, total_size, spec)?;
    let prefix_cap =
        usize::try_from(total_size.min(u64::try_from(AVS3_PREFIX_SCAN_BYTES).unwrap()))
            .map_err(|_| MuxError::LayoutOverflow("transport-stream AVS3 prefix read length"))?;
    let mut sequence_header = None::<ParsedAvs3SequenceHeader>;
    let mut samples = Vec::with_capacity(sample_bounds.len());

    for (sample_offset, data_size) in sample_bounds {
        let prefix_len = prefix_cap.min(usize::try_from(data_size).unwrap_or(prefix_cap));
        let mut prefix = vec![0_u8; prefix_len];
        read_segmented_bytes_sync(
            file,
            segments,
            total_size,
            sample_offset,
            &mut prefix,
            spec,
            "transport-stream AVS3 sample prefix is truncated",
        )?;
        if sequence_header.is_none()
            && let Some(sequence_start) =
                find_prefixed_start_code(&prefix, AVS3_SEQUENCE_HEADER_PREFIX)
        {
            sequence_header = Some(parse_avs3_sequence_header(spec, &prefix[sequence_start..])?);
        }
        samples.push(StagedSample {
            data_offset: sample_offset,
            data_size,
            duration: 0,
            composition_time_offset: 0,
            is_sync_sample: classify_transport_avs3_sample(spec, &prefix)?,
        });
    }

    finalize_transport_avs3_track(spec, config, sequence_header, samples)
}

#[cfg(feature = "async")]
pub(in crate::mux) async fn scan_transport_avs3_segmented_async(
    file: &mut TokioFile,
    segments: &[SegmentedMuxSourceSegment],
    total_size: u64,
    sample_offsets: &[u64],
    avs3_config: &[u8],
    spec: &str,
) -> Result<ParsedTransportAvs3Track, MuxError> {
    let config = parse_transport_avs3_config_bytes(avs3_config, spec)?;
    let sample_bounds = build_transport_avs3_sample_bounds(sample_offsets, total_size, spec)?;
    let prefix_cap =
        usize::try_from(total_size.min(u64::try_from(AVS3_PREFIX_SCAN_BYTES).unwrap()))
            .map_err(|_| MuxError::LayoutOverflow("transport-stream AVS3 prefix read length"))?;
    let mut sequence_header = None::<ParsedAvs3SequenceHeader>;
    let mut samples = Vec::with_capacity(sample_bounds.len());

    for (sample_offset, data_size) in sample_bounds {
        let prefix_len = prefix_cap.min(usize::try_from(data_size).unwrap_or(prefix_cap));
        let mut prefix = vec![0_u8; prefix_len];
        read_segmented_bytes_async(
            file,
            segments,
            total_size,
            sample_offset,
            &mut prefix,
            spec,
            "transport-stream AVS3 sample prefix is truncated",
        )
        .await?;
        if sequence_header.is_none()
            && let Some(sequence_start) =
                find_prefixed_start_code(&prefix, AVS3_SEQUENCE_HEADER_PREFIX)
        {
            sequence_header = Some(parse_avs3_sequence_header(spec, &prefix[sequence_start..])?);
        }
        samples.push(StagedSample {
            data_offset: sample_offset,
            data_size,
            duration: 0,
            composition_time_offset: 0,
            is_sync_sample: classify_transport_avs3_sample(spec, &prefix)?,
        });
    }

    finalize_transport_avs3_track(spec, config, sequence_header, samples)
}

fn finalize_transport_avs3_track(
    spec: &str,
    config: Av3c,
    sequence_header: Option<ParsedAvs3SequenceHeader>,
    mut samples: Vec<StagedSample>,
) -> Result<ParsedTransportAvs3Track, MuxError> {
    let sequence_header = sequence_header.ok_or(MuxError::UnsupportedTrackImport {
        spec: spec.to_string(),
        message:
            "transport-stream AVS3 carriage did not expose a sequence header on the native direct-ingest path"
                .to_string(),
    })?;
    if !sequence_header.low_delay {
        return Err(MuxError::UnsupportedTrackImport {
            spec: spec.to_string(),
            message:
                "transport-stream AVS3 carriage with reordered presentation is not supported on the native direct-ingest path yet"
                    .to_string(),
        });
    }
    for sample in &mut samples {
        sample.duration = sequence_header.sample_duration;
    }
    let config_box = super::super::mp4::encode_typed_box(&config, &[])?;
    let btrt_box = super::super::mp4::encode_typed_box(
        &build_btrt_from_sample_sizes(
            samples
                .iter()
                .map(|sample| (sample.data_size, sample.duration)),
            sequence_header.timescale,
        )?,
        &[],
    )?;
    Ok(ParsedTransportAvs3Track {
        timescale: sequence_header.timescale,
        sample_entry_box: build_visual_sample_entry_box(
            SAMPLE_ENTRY_AVS3,
            0,
            0,
            &[config_box, btrt_box],
        )?,
        samples,
    })
}

fn build_transport_avs3_sample_bounds(
    sample_offsets: &[u64],
    total_size: u64,
    spec: &str,
) -> Result<Vec<(u64, u32)>, MuxError> {
    if sample_offsets.is_empty() {
        return Err(MuxError::UnsupportedTrackImport {
            spec: spec.to_string(),
            message: "transport-stream AVS3 carriage did not contain any PES payload units"
                .to_string(),
        });
    }
    let mut bounds = Vec::with_capacity(sample_offsets.len());
    for (index, &sample_offset) in sample_offsets.iter().enumerate() {
        let next_offset = sample_offsets.get(index + 1).copied().unwrap_or(total_size);
        if next_offset <= sample_offset {
            return Err(MuxError::UnsupportedTrackImport {
                spec: spec.to_string(),
                message: "transport-stream AVS3 PES payload units must advance monotonically"
                    .to_string(),
            });
        }
        bounds.push((
            sample_offset,
            u32::try_from(next_offset - sample_offset)
                .map_err(|_| MuxError::LayoutOverflow("transport-stream AVS3 sample size"))?,
        ));
    }
    Ok(bounds)
}

fn parse_transport_avs3_config_bytes(config_bytes: &[u8], spec: &str) -> Result<Av3c, MuxError> {
    if config_bytes.len() != AVS3_DESCRIPTOR_CONFIG_SIZE {
        return Err(MuxError::UnsupportedTrackImport {
            spec: spec.to_string(),
            message:
                "transport-stream AVS3 registration descriptor did not carry the expected 10-byte decoder configuration payload"
                    .to_string(),
        });
    }
    let sequence_header_length = u16::from_be_bytes([config_bytes[1], config_bytes[2]]);
    let expected_size =
        usize::from(sequence_header_length)
            .checked_add(4)
            .ok_or(MuxError::LayoutOverflow(
                "transport-stream AVS3 decoder config size",
            ))?;
    if expected_size != config_bytes.len() {
        return Err(MuxError::UnsupportedTrackImport {
            spec: spec.to_string(),
            message:
                "transport-stream AVS3 decoder configuration length does not match its carried payload"
                    .to_string(),
        });
    }
    let sequence_header_end = 3 + usize::from(sequence_header_length);
    let sequence_header = config_bytes[3..sequence_header_end].to_vec();
    if !sequence_header.starts_with(&AVS3_SEQUENCE_HEADER_PREFIX) {
        return Err(MuxError::UnsupportedTrackImport {
            spec: spec.to_string(),
            message:
                "transport-stream AVS3 decoder configuration did not begin with a sequence-header prefix"
                    .to_string(),
        });
    }
    Ok(Av3c {
        configuration_version: config_bytes[0],
        sequence_header_length,
        sequence_header,
        library_dependency_idc: config_bytes[config_bytes.len() - 1] & 0x03,
    })
}

fn parse_avs3_sequence_header(
    spec: &str,
    bytes: &[u8],
) -> Result<ParsedAvs3SequenceHeader, MuxError> {
    if bytes.len() < 17 || !bytes.starts_with(&AVS3_SEQUENCE_HEADER_PREFIX) {
        return Err(MuxError::UnsupportedTrackImport {
            spec: spec.to_string(),
            message: "transport-stream AVS3 sequence header is truncated".to_string(),
        });
    }
    let mut reader = BitReader::new(Cursor::new(&bytes[4..]));
    let _profile = read_bits_u8_labeled(&mut reader, 8, spec, "transport AVS3 profile")?;
    let _level = read_bits_u8_labeled(&mut reader, 8, spec, "transport AVS3 level")?;
    let _progressive = read_bit_labeled(&mut reader, spec, "transport AVS3 progressive flag")?;
    let _field = read_bit_labeled(&mut reader, spec, "transport AVS3 field flag")?;
    let _library_stream =
        read_bits_u8_labeled(&mut reader, 2, spec, "transport AVS3 library-stream flag")?;
    skip_bits_labeled(&mut reader, 1, spec, "transport AVS3 reserved bit")?;
    let width = read_bits_u16_labeled(&mut reader, 14, spec, "transport AVS3 width")?;
    skip_bits_labeled(&mut reader, 1, spec, "transport AVS3 reserved width bit")?;
    let height = read_bits_u16_labeled(&mut reader, 14, spec, "transport AVS3 height")?;
    skip_bits_labeled(&mut reader, 2, spec, "transport AVS3 chroma format")?;
    skip_bits_labeled(&mut reader, 3, spec, "transport AVS3 sample precision")?;
    skip_bits_labeled(&mut reader, 1, spec, "transport AVS3 reserved aspect bit")?;
    skip_bits_labeled(&mut reader, 4, spec, "transport AVS3 aspect ratio")?;
    let frame_rate_code =
        read_bits_u8_labeled(&mut reader, 4, spec, "transport AVS3 frame-rate code")?;
    skip_bits_labeled(&mut reader, 1, spec, "transport AVS3 reserved bitrate bit")?;
    skip_bits_labeled(&mut reader, 18, spec, "transport AVS3 bitrate low bits")?;
    skip_bits_labeled(
        &mut reader,
        1,
        spec,
        "transport AVS3 reserved high-bitrate bit",
    )?;
    skip_bits_labeled(&mut reader, 12, spec, "transport AVS3 bitrate high bits")?;
    let low_delay = read_bit_labeled(&mut reader, spec, "transport AVS3 low-delay flag")?;
    if width == 0 || height == 0 {
        return Err(MuxError::UnsupportedTrackImport {
            spec: spec.to_string(),
            message: "transport-stream AVS3 sequence header declared a zero coded picture size"
                .to_string(),
        });
    }
    let (timescale, sample_duration) = avs3_frame_rate_from_code(frame_rate_code, spec)?;
    Ok(ParsedAvs3SequenceHeader {
        timescale,
        sample_duration,
        low_delay,
    })
}

fn avs3_frame_rate_from_code(frame_rate_code: u8, spec: &str) -> Result<(u32, u32), MuxError> {
    match frame_rate_code {
        0x01 => Ok((24_000, 1_001)),
        0x02 => Ok((24, 1)),
        0x03 => Ok((25, 1)),
        0x04 => Ok((30_000, 1_001)),
        0x05 => Ok((30, 1)),
        0x06 => Ok((50, 1)),
        0x07 => Ok((60_000, 1_001)),
        0x08 => Ok((60, 1)),
        0x09 => Ok((100, 1)),
        0x0A => Ok((120, 1)),
        0x0B => Ok((200, 1)),
        0x0C => Ok((240, 1)),
        0x0D => Ok((300, 1)),
        _ => Err(MuxError::UnsupportedTrackImport {
            spec: spec.to_string(),
            message: format!(
                "transport-stream AVS3 sequence header used unsupported frame-rate code 0x{frame_rate_code:02X}"
            ),
        }),
    }
}

fn classify_transport_avs3_sample(spec: &str, bytes: &[u8]) -> Result<bool, MuxError> {
    if find_prefixed_start_code(bytes, AVS3_INTRA_PICTURE_PREFIX).is_some() {
        return Ok(true);
    }
    if find_prefixed_start_code(bytes, AVS3_INTER_PICTURE_PREFIX).is_some() {
        return Ok(false);
    }
    Err(MuxError::UnsupportedTrackImport {
        spec: spec.to_string(),
        message: "transport-stream AVS3 sample did not contain a supported picture-start prefix"
            .to_string(),
    })
}

fn find_prefixed_start_code(bytes: &[u8], needle: [u8; 4]) -> Option<usize> {
    bytes.windows(4).position(|window| window == needle)
}