diff --git a/README.md b/README.md index 45361e66e..ab45c1a99 100644 --- a/README.md +++ b/README.md @@ -171,8 +171,8 @@ craft- = { git = "https://github.com/storytold/craft-libs", rev = " = file.read_unit(&bytes, v, 0)?; // the first access unit +``` + +- **Transport streams**: 188-byte packets, 192-byte BDAV packets (sync byte after a 4-byte + arrival timestamp) and 204-byte packets; the packet size is sniffed from five consecutive sync + bytes. Lost sync resynchronises on the next aligned run. PSI sections are reassembled across + packets and checked with the MPEG CRC-32; the first program of the PAT is demuxed. Continuity + counter errors, CRC errors and resyncs are reported as warnings. +- **Program streams**: MPEG-2 and MPEG-1 packs, system headers, the program stream map (stream + types for elementary stream ids), padding, private stream 2 (DVD navigation) skipped, private + stream 1 split by sub-stream id. Garbage resynchronises on the next pack or PES start code. +- **Streams**: MPEG-1/2 video, H.264, HEVC, MPEG audio, AAC (ADTS, LATM/LOAS), AC-3, E-AC-3, + DTS, TrueHD, Blu-ray LPCM, DVD LPCM (format from the sub-stream header), PGS / DVB / teletext + subtitles; ISO 639 language from the PMT. +- **Access units**, found while the file is scanned (one pass, streaming; only a few bytes of + each stream are buffered): + - MPEG-1/2 video: a coded frame — a frame picture, or two field pictures (paired by + picture_structure) — with the sequence / GOP headers before it; picture type, temporal + reference, field structure, top_field_first, repeat_first_field, progressive_frame, closed + GOP / broken link. A PTS belongs to the unit whose picture start code is in that PES. + - H.264 / HEVC: PES packets with a timestamp (a PES without one continues the previous unit); + random access from IDR / IRAP NAL types; disposable when every slice is a non-reference. + - Audio: one frame per unit from the frame headers (sync re-hunted after garbage; a header + found while hunting is confirmed by the next one); a PTS belongs to the first frame starting + in the PES. + - LPCM and everything else: one PES packet per unit. +- **Timestamps**: 33-bit PTS/DTS unwrapped across the wrap (and back, for B pictures shown + before a wrap that came earlier); PCR / SCR range. +- **Reading**: a unit records only the position of its first packet and its offset in that + packet's payload; `read_unit` re-reads from there (the index costs a few dozen bytes per + unit, not per packet). + +`frame_bytes`, `frame_info` and `sei_has_recovery_point` are public: other containers that store +MPEG audio or AC-3 frames in chunks that do not follow them (AVI) find the frames again with them. + +`Codec::codec_id()` and `Stream::codec_params()` give the crate-wide `CodecId` / `CodecParams`. +MPEG-1 video reports `Mpeg2Video`; MPEG audio reports `Mp2` whatever its layer (the layer is in +each frame header: `frame_info(..).variant`); ADTS and LATM AAC both report `Aac`; DTS, TrueHD and +subtitles have no identifier (`Unknown`). The streams carry their configuration in band, so +`extradata` is empty; DVD LPCM fills `sample_rate`, `channels` and `pcm`. + +Limits, each with a test: 2²³ access units per stream (the rest are not indexed, with a warning), +128 MiB per unit read (`MAX_UNIT_BYTES`), a PES header given up after 1024 bytes, 16 stuffing bytes +in an MPEG-1 packet header, a pack header sniffed behind at most 64 zero bytes, and E-AC-3 frames of +at least 8 bytes. + +### Specifications + +| Document | Edition | Used for | +|---|---|---| +| ITU-T H.222.0 \| ISO/IEC 13818-1 | 2000 (second edition) | TS packets, adaptation fields, PCR, PSI (PAT / PMT, CRC-32 of Annex A), PES headers, PTS/DTS, program streams (pack, system header, program stream map), stream types | +| ISO/IEC 11172-1 | 1993 | MPEG-1 system streams (pack header, PES header stuffing / STD buffer / timestamps) | +| Blu-ray / AVCHD BDAV MPEG-2 TS layout (192-byte source packets with a 4-byte TP_extra_header; registration `HDMV`; stream types 0x80-0x92) | public descriptions | `.m2ts` / `.mts` | +| DVD-Video private stream 1 sub-stream layout (AC-3 / DTS 0x80-0x8F with a 3-byte header, LPCM 0xA0-0xAF with a 6-byte header, sub-pictures 0x20-0x3F) | public descriptions | `.vob` | +| ETSI EN 300 468 (descriptors 0x56, 0x59, 0x6A, 0x7A, 0x7B), ATSC A/52 Annex A (AC-3 in TS) | — | identifying private PES streams | +| ISO/IEC 11172-3, 13818-3 (frame headers), ISO/IEC 13818-7 (ADTS), 14496-3 (LOAS), ATSC A/52 (syncinfo, E-AC-3 frmsiz) | — | audio frame lengths for splitting | + +### Gaps + +- One program per TS (the first in the PAT); no program switching, no section-based + metadata (SDT / EIT), no scrambled streams. +- H.264 / HEVC access units follow PES packets: streams that put several pictures in one PES + packet (rare in practice; not seen in broadcast, AVCHD or FFmpeg output) would be indexed as + one unit. +- DVD navigation (IFO / NAV packets), multiple VOB concatenation and angle interleaving are not + handled; each file opens on its own. + +Differences from `filmcraft-mpegts`: `ByteSource` is `craft_media::ByteSource`; +`Error` is written by hand, without `thiserror`, and is `#[non_exhaustive]`; +`Codec::codec_id()` and `Stream::codec_params()` are new. + +Tests: a small TS muxer and PS writer (every access unit reads back byte-identical with its PTS, +key and disposable flags), 33-bit unwrap, truncation and mutation, streams of every codec filled +with random start codes and sync words, damaged PSI sections with valid CRCs, and the frame-header +helpers on arbitrary bytes. `tests/oracle_mpegts.rs` compares ten ffmpeg-made files (MPEG-2, H.264 +and HEVC video; MP2, ADTS and LATM AAC, AC-3, E-AC-3, Blu-ray and DVD LPCM; TS, M2TS, VOB and +MPEG-1 system streams) with ffprobe's packets: unit counts, sizes, timestamps and key flags. ## `mxf` @@ -664,7 +757,7 @@ sound. Always built, with no dependencies, so codec crates share them: - `Container`: the file format (`Wav`, `Aiff`, `Flac`, `Mp4`, `QuickTime`, `Matroska`, `WebM`, - `Ogg`, `Caf`, `Adts`, `MpegAudio`, `Ac3`, `Avi`, `Mxf`), with `name()`. + `Ogg`, `Caf`, `Adts`, `MpegAudio`, `Ac3`, `Avi`, `Mxf`, `MpegTs`, `MpegPs`), with `name()`. - `CodecId`: audio and video codecs (`Aac`, `HeAac`, `Opus`, `Eac3`, `Speex`, …, `H264`, `Hevc`, `Av1`, `ProRes`, `Dnxhd`, `Apv`, `Theora`, …, `Unknown`), with `name()`, `is_audio()`, `is_video()`. @@ -686,7 +779,7 @@ All are `#[non_exhaustive]`: new formats and fields are minor changes. `craft_media::ByteSource` (`len`, `read_at(offset, buf)`) is what every demuxer reads from. It is always built, so a crate that uses only `bitstream` can name it too. `isobmff::ByteSource`, -`matroska::ByteSource`, `ogg::ByteSource`, `avi::ByteSource` and `mxf::ByteSource` are the same trait. It is implemented for `[u8]`, +`matroska::ByteSource`, `ogg::ByteSource`, `avi::ByteSource`, `mxf::ByteSource` and `mpegts::ByteSource` are the same trait. It is implemented for `[u8]`, `Vec`, `&T`, `Arc`, `Box` and `std::fs::File` (not on wasm). `ReadSeekSource` adapts any `Read + Seek`; it is `Send + Sync` when the reader is `Send`. Implement the trait once for a custom source (a web `Blob`, a network range reader) to use every container with it. diff --git a/crates/media/src/id.rs b/crates/media/src/id.rs index 7862546f3..00b63d004 100644 --- a/crates/media/src/id.rs +++ b/crates/media/src/id.rs @@ -40,6 +40,10 @@ pub enum Container { Avi, /// MXF (SMPTE ST 377-1): OP1a, OP-Atom. Mxf, + /// MPEG-2 transport stream (TS, M2TS). + MpegTs, + /// MPEG-1/2 program stream (MPG, VOB). + MpegPs, } impl Container { @@ -60,6 +64,8 @@ impl Container { Container::Ac3 => "AC-3", Container::Avi => "AVI", Container::Mxf => "MXF", + Container::MpegTs => "MPEG-TS", + Container::MpegPs => "MPEG-PS", } } } diff --git a/crates/media/src/lib.rs b/crates/media/src/lib.rs index eaaea04ed..a9c812318 100644 --- a/crates/media/src/lib.rs +++ b/crates/media/src/lib.rs @@ -10,10 +10,12 @@ //! - `ogg` (feature `ogg`): a clean-room Ogg demuxer with Ogg Opus and Ogg Vorbis stream mapping. //! - `avi` (feature `avi`): a clean-room AVI (RIFF) and OpenDML demuxer. //! - `mxf` (feature `mxf`, enables `chrono`): a clean-room MXF demuxer and writer (OP1a, OP-Atom). +//! - `mpegts` (feature `mpegts`): a clean-room MPEG-2 Systems demuxer (transport streams and +//! program streams). //! - [`Container`] and [`CodecId`] (no feature): the identifiers the codec crates report a file's //! container and codec with. //! -//! `bitstream` and `isobmff` are on by default; `matroska`, `ogg`, `avi` and `mxf` are opt-in. A crate that needs +//! `bitstream` and `isobmff` are on by default; `matroska`, `ogg`, `avi`, `mxf` and `mpegts` are opt-in. A crate that needs //! only bit I/O uses `default-features = false, features = ["bitstream"]`. #[cfg(feature = "avi")] @@ -25,6 +27,8 @@ mod id; pub mod isobmff; #[cfg(feature = "matroska")] pub mod matroska; +#[cfg(feature = "mpegts")] +pub mod mpegts; #[cfg(feature = "mxf")] pub mod mxf; #[cfg(feature = "ogg")] diff --git a/crates/media/src/mpegts/es.rs b/crates/media/src/mpegts/es.rs new file mode 100644 index 000000000..8568c7cba --- /dev/null +++ b/crates/media/src/mpegts/es.rs @@ -0,0 +1,902 @@ +//! Splitting elementary streams into access units while the file is scanned. +//! +//! The scanners feed each stream's payload bytes chunk by chunk (a TS packet's payload, a PS +//! packet's payload). A chunk is identified by the position the reader restarts from (`pos`) and +//! the elementary-stream offset of its first byte. Splitting needs only a few bytes beyond a +//! boundary, so only the most recent chunks are remembered. + +use std::collections::VecDeque; + +use super::{Codec, PictureInfo, Unit}; + +/// How a stream is cut into units. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Mode { + /// MPEG-1/2 video: one coded frame (a frame picture or two field pictures) per unit. + MpegVideo, + /// H.264 / HEVC: one access unit per unit (the first of a PES packet keeps its timestamps); + /// NAL types give the flags. + Nal { hevc: bool }, + /// Framed audio with self-describing frame headers. + Audio(AudioSync), + /// One PES packet per unit. + Pes, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum AudioSync { + Mpa, + Adts, + Loas, + Ac3, +} + +impl Mode { + fn of(codec: &Codec) -> Mode { + match codec { + Codec::Mpeg1Video | Codec::Mpeg2Video => Mode::MpegVideo, + Codec::H264 => Mode::Nal { hevc: false }, + Codec::Hevc => Mode::Nal { hevc: true }, + Codec::MpegAudio => Mode::Audio(AudioSync::Mpa), + Codec::AacAdts => Mode::Audio(AudioSync::Adts), + Codec::AacLatm => Mode::Audio(AudioSync::Loas), + Codec::Ac3 | Codec::Eac3 => Mode::Audio(AudioSync::Ac3), + _ => Mode::Pes, + } + } +} + +#[derive(Clone, Copy)] +struct PesMark { + es_start: u64, + pts: Option, + dts: Option, + used: bool, +} + +/// A unit whose end is not known yet. +#[derive(Clone, Copy)] +struct Open { + start: u64, + pos: u64, + offset: u32, + pts: Option, + dts: Option, + key: bool, + disposable: bool, + /// A coded slice has been seen (H.264 / HEVC). + vcl: bool, + picture: Option, +} + +/// MPEG video picture waiting for its coding extension (to tell fields from frames). +#[derive(Clone, Copy)] +struct PendingPicture { + /// Start of the unit if the picture begins one (its first sequence / GOP header). + unit_start: u64, + /// The picture start code (timestamps belong to the unit whose picture starts in the PES). + at: u64, + info: PictureInfo, +} + +pub(crate) struct Splitter { + mode: Mode, + /// Elementary-stream bytes fed so far. + es_pos: u64, + /// Recent chunks: (reader position, ES offset of the first byte). + chunks: VecDeque<(u64, u64)>, + pes: VecDeque, + /// Unconsumed bytes (start codes / frame headers spanning chunks), starting at `buf_start`. + buf: Vec, + buf_start: u64, + open: Option, + pub units: Vec, + pub skipped: u64, + /// Units beyond [`MAX_UNITS`] were dropped. + pub truncated: bool, + max_units: usize, + // MPEG video state + header_start: Option, + seq_in_unit: bool, + gop: Option<(bool, bool)>, + pending: Option, + open_field: Option, + // audio state + next_frame: Option, +} + +/// AC-3 frame sizes in 16-bit words for 48 / 44.1 / 32 kHz (A/52 Table 5.18), by +/// frmsizecod / 2; 44.1 kHz adds a word for odd frmsizecod. +const AC3_KBPS: [u32; 19] = [32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 384, 448, 512, 576, 640]; + +pub(crate) fn ac3_frame_bytes(fscod: u8, frmsizecod: u8) -> Option { + let kbps = *AC3_KBPS.get((frmsizecod / 2) as usize)?; + let words = match fscod { + 0 => kbps * 2, + 1 => kbps * 96_000 / 44_100 + (frmsizecod & 1) as u32, + 2 => kbps * 3, + _ => return None, + }; + Some(words as usize * 2) +} + +const MPA_KBPS: [[u16; 15]; 5] = [ + [0, 32, 64, 96, 128, 160, 192, 224, 256, 288, 320, 352, 384, 416, 448], // MPEG-1 layer I + [0, 32, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 384], // MPEG-1 layer II + [0, 32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320], // MPEG-1 layer III + [0, 32, 48, 56, 64, 80, 96, 112, 128, 144, 160, 176, 192, 224, 256], // MPEG-2/2.5 layer I + [0, 8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160], // MPEG-2/2.5 layers II, III +]; + +/// MPEG audio frame header (ISO/IEC 11172-3 §2.4.2.3, 13818-3 LSF): (frame bytes, sample rate, +/// samples per frame). +pub(crate) fn mpa_header(b: &[u8]) -> Option<(usize, u32, u32)> { + if b.len() < 4 || b[0] != 0xFF || b[1] & 0xE0 != 0xE0 { + return None; + } + let version = (b[1] >> 3) & 3; // 0: 2.5, 2: MPEG-2, 3: MPEG-1 + let layer = (b[1] >> 1) & 3; // 3: I, 2: II, 1: III + let br = (b[2] >> 4) as usize; + let sr = (b[2] >> 2) & 3; + let pad = ((b[2] >> 1) & 1) as usize; + if version == 1 || layer == 0 || br == 0 || br == 15 || sr == 3 { + return None; + } + let rate = [44_100u32, 48_000, 32_000][sr as usize] + >> match version { + 3 => 0, + 2 => 1, + _ => 2, + }; + let table = match (version == 3, layer) { + (true, 3) => 0, + (true, 2) => 1, + (true, _) => 2, + (false, 3) => 3, + (false, _) => 4, + }; + let kbps = MPA_KBPS[table][br] as usize; + let (bytes, samples) = match layer { + 3 => ((12 * kbps * 1000 / rate as usize + pad) * 4, 384), + 2 => (144 * kbps * 1000 / rate as usize + pad, 1152), + _ if version == 3 => (144 * kbps * 1000 / rate as usize + pad, 1152), + _ => (72 * kbps * 1000 / rate as usize + pad, 576), + }; + (bytes >= 4).then_some((bytes, rate, samples)) +} + +/// What an audio frame header says. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct FrameInfo { + /// Samples per second. + pub sample_rate: u32, + /// Channel count. + pub channels: u32, + /// Samples per channel in the frame. + pub samples: u32, + /// MPEG audio layer (1-3); AC-3 bsid; ADTS object type (profile + 1). + pub variant: u8, +} + +const ADTS_RATES: [u32; 13] = [96_000, 88_200, 64_000, 48_000, 44_100, 32_000, 24_000, 22_050, 16_000, 12_000, 11_025, 8_000, 7_350]; + +/// Bytes of the audio frame whose header starts `b` (MPEG audio, ADTS, LATM, AC-3, E-AC-3), or +/// `None` without a valid header there. For containers that store frames in chunks that need not +/// follow them (AVI): the frames are found again with it. +pub fn frame_bytes(codec: &Codec, b: &[u8]) -> Option { + let sync = match codec { + Codec::MpegAudio => AudioSync::Mpa, + Codec::AacAdts => AudioSync::Adts, + Codec::AacLatm => AudioSync::Loas, + Codec::Ac3 | Codec::Eac3 => AudioSync::Ac3, + _ => return None, + }; + frame_len(sync, b) +} + +/// Sample rate, channels and samples of the audio frame whose header starts `b` (the codecs of +/// [`frame_bytes`] except LATM), or `None` without a valid header there. +pub fn frame_info(codec: &Codec, b: &[u8]) -> Option { + match codec { + Codec::MpegAudio => { + let (_, sample_rate, samples) = mpa_header(b)?; + let layer = 4 - ((b[1] >> 1) & 3); + Some(FrameInfo { sample_rate, channels: if b[3] >> 6 == 3 { 1 } else { 2 }, samples, variant: layer }) + } + Codec::AacAdts => { + frame_len(AudioSync::Adts, b)?; + let sf = ((b[2] >> 2) & 15) as usize; + let ch = (((b[2] & 1) << 2) | (b[3] >> 6)) as u32; + Some(FrameInfo { + sample_rate: *ADTS_RATES.get(sf)?, + channels: if ch == 0 { + 2 + } else if ch == 7 { + 8 + } else { + ch + }, + samples: 1024 * ((b[6] & 3) as u32 + 1), + variant: (b[2] >> 6) + 1, + }) + } + Codec::Ac3 | Codec::Eac3 => { + frame_len(AudioSync::Ac3, b)?; + let bsid = b[5] >> 3; + if bsid > 10 { + // E-AC-3: fscod(2) numblkscod(2) acmod(3) lfeon(1) + let fscod = b[4] >> 6; + let (rate, blocks) = if fscod == 3 { + // fscod2 3 is reserved (no sample rate) + (*[24_000, 22_050, 16_000].get(((b[4] >> 4) & 3) as usize)?, 6) + } else { + ([48_000, 44_100, 32_000][fscod as usize], [1, 2, 3, 6][((b[4] >> 4) & 3) as usize]) + }; + let acmod = (b[4] >> 1) & 7; + return Some(FrameInfo { + sample_rate: rate, + channels: AC3_CHANNELS[acmod as usize] + (b[4] & 1) as u32, + samples: 256 * blocks, + variant: bsid, + }); + } + let fscod = b[4] >> 6; + let rate = *[48_000u32, 44_100, 32_000].get(fscod as usize)?; + // bsi: bsid(5) bsmod(3) acmod(3) [cmixlev(2)] [surmixlev(2)] [dsurmod(2)] lfeon(1) + let b6 = *b.get(6)?; + let acmod = b6 >> 5; + let mut bit = 3; // bits of byte 6 consumed + if acmod & 1 != 0 && acmod != 1 { + bit += 2; + } + if acmod & 4 != 0 { + bit += 2; + } + if acmod == 2 { + bit += 2; + } + let word = u16::from_be_bytes([b6, *b.get(7)?]); + let lfe = (word >> (15 - bit)) & 1; + Some(FrameInfo { sample_rate: rate, channels: AC3_CHANNELS[acmod as usize] + lfe as u32, samples: 1536, variant: bsid }) + } + _ => None, + } +} + +/// Whether the syncframe at the start of `b` is E-AC-3 (bsid 11-16) other than independent +/// substream 0: a dependent substream (strmtyp 1) or another program (A/52 §E2.3.1.1-2). +fn eac3_joins_unit(b: &[u8]) -> bool { + match b.get(2..6) { + Some(&[b2, _, _, b5]) => (11..=16).contains(&(b5 >> 3)) && (b2 >> 6 == 1 || (b2 >> 3) & 7 != 0), + _ => false, + } +} + +/// Bytes of an SEI NAL scanned for a recovery point message; a longer NAL is judged by its start. +const SEI_WINDOW: usize = 2048; + +/// Whether an H.264 SEI NAL (`nal` is the NAL header byte and the payload, as in the stream) +/// holds a recovery point message (payloadType 6, H.264 D.1.7) with recovery_frame_cnt 0. +pub fn sei_has_recovery_point(nal: &[u8]) -> bool { + // undo emulation prevention so that payload sizes are in RBSP bytes + let mut b = Vec::with_capacity(nal.len()); + let mut zeros = 0; + for &x in nal.iter().skip(1) { + if zeros >= 2 && x == 3 { + zeros = 0; + continue; + } + zeros = if x == 0 { zeros + 1 } else { 0 }; + b.push(x); + } + // payloadType / payloadSize: bytes of 0xFF followed by the last byte + let read = |i: &mut usize| { + let mut v = 0usize; + while let Some(&x) = b.get(*i) { + *i += 1; + v = v.saturating_add(x as usize); + if x != 0xFF { + break; + } + } + v + }; + let mut i = 0; + // each message takes at least two bytes, so the loop ends at the end of the NAL + while i + 2 <= b.len() { + let ty = read(&mut i); + let size = read(&mut i); + // recovery_frame_cnt is ue(v): a leading 1 bit is the value 0, i.e. decoding from here is + // correct at once. A count above 0 (gradual intra refresh) is not a safe seek point. + if ty == 6 && b.get(i).is_some_and(|x| x & 0x80 != 0) { + return true; + } + i = i.saturating_add(size); + } + false +} + +/// Shortest E-AC-3 syncframe accepted. A real syncframe holds the 6 header bytes read here plus +/// bsi and audio blocks (the smallest real ones are tens of bytes), so shorter "frames" +/// (`0B 77 00 01` is 4 bytes) are noise; they would otherwise turn every 4 input bytes into a +/// 64-byte `Unit`. +const EAC3_MIN_FRAME: usize = 8; + +/// Most units one stream may index. A `Unit` is 64 bytes, so this bounds a stream's table to +/// 512 MiB. It is 47 hours of 48 kHz AAC (the densest audio, 46.9 units/s), 9 hours of 240 fps +/// video: far beyond any real file, while a hostile file (a unit per 7 input bytes) cannot grow +/// the table without limit. +pub(crate) const MAX_UNITS: usize = 1 << 23; + +/// Full-bandwidth channels per AC-3 audio coding mode (A/52 Table 5.8). +const AC3_CHANNELS: [u32; 8] = [2, 1, 2, 3, 3, 4, 4, 5]; + +/// Frame length at the start of `b` for the given sync type (`None`: no valid header; `Some(0)` +/// is never returned). +fn frame_len(sync: AudioSync, b: &[u8]) -> Option { + match sync { + AudioSync::Mpa => mpa_header(b).map(|h| h.0), + AudioSync::Adts => { + if b.len() < 7 || b[0] != 0xFF || b[1] & 0xF6 != 0xF0 { + return None; + } + let len = ((b[3] as usize & 3) << 11) | (b[4] as usize) << 3 | (b[5] as usize) >> 5; + (len >= 7).then_some(len) + } + AudioSync::Loas => { + if b.len() < 3 || b[0] != 0x56 || b[1] & 0xE0 != 0xE0 { + return None; + } + let len = ((b[1] as usize & 0x1F) << 8) | b[2] as usize; + (len > 0).then_some(len + 3) + } + AudioSync::Ac3 => { + if b.len() < 6 || b[0] != 0x0B || b[1] != 0x77 { + return None; + } + let bsid = b[5] >> 3; + if bsid <= 10 { + ac3_frame_bytes(b[4] >> 6, b[4] & 0x3F) + } else if bsid <= 16 { + // E-AC-3: frmsiz + let len = ((((b[2] & 7) as usize) << 8) | b[3] as usize) * 2 + 2; + (len >= EAC3_MIN_FRAME).then_some(len) + } else { + None + } + } + } +} + +impl Splitter { + pub fn new(codec: &Codec) -> Splitter { + Splitter { + mode: Mode::of(codec), + es_pos: 0, + chunks: VecDeque::new(), + pes: VecDeque::new(), + buf: Vec::new(), + buf_start: 0, + open: None, + units: Vec::new(), + skipped: 0, + truncated: false, + max_units: MAX_UNITS, + header_start: None, + seq_in_unit: false, + gop: None, + pending: None, + open_field: None, + next_frame: None, + } + } + + /// A PES packet starts (its payload follows in the next chunks). + pub fn pes_start(&mut self, pts: Option, dts: Option) { + self.pes.push_back(PesMark { es_start: self.es_pos, pts, dts, used: false }); + while self.pes.len() > 64 { + self.pes.pop_front(); + } + match self.mode { + Mode::Pes => self.begin_pes_unit(true), + Mode::Nal { .. } => self.begin_pes_unit(pts.is_some() || dts.is_some()), + _ => {} + } + } + + /// PES-delimited units: a packet with a timestamp (or any packet in plain PES mode) starts a + /// new unit; one without continues the previous unit. + fn begin_pes_unit(&mut self, new: bool) { + if !new && self.open.is_some() { + return; + } + let at = self.es_pos; + self.close(at); + // the chunk of this position arrives with the next feed: resolved there + let mut o = Open { start: at, pos: u64::MAX, offset: 0, pts: None, dts: None, key: false, disposable: false, vcl: false, picture: None }; + if let Some((pts, dts)) = self.timestamps_at(at) { + o.pts = pts; + o.dts = dts; + } + if self.mode == Mode::Pes { + o.key = true; + } + self.open = Some(o); + } + + /// The unused timestamps of the PES packet holding `at` (consumed). + fn timestamps_at(&mut self, at: u64) -> Option<(Option, Option)> { + let m = self.pes.iter_mut().rev().find(|m| m.es_start <= at)?; + if m.used || (m.pts.is_none() && m.dts.is_none()) { + return None; + } + m.used = true; + Some((m.pts, m.dts)) + } + + /// Reader position and offset of elementary-stream byte `at`. + fn locate(&self, at: u64) -> Option<(u64, u32)> { + let c = self.chunks.iter().rev().find(|c| c.1 <= at)?; + Some((c.0, (at - c.1) as u32)) + } + + /// End the open unit at `end`. + fn close(&mut self, end: u64) { + let Some(o) = self.open.take() else { return }; + if end <= o.start || o.pos == u64::MAX { + return; + } + if self.units.len() >= self.max_units { + self.truncated = true; + return; + } + let size = (end - o.start).min(u32::MAX as u64) as u32; + self.units.push(Unit { + pts: o.pts, + dts: o.dts, + key: o.key, + disposable: o.disposable, + size, + pos: o.pos, + offset: o.offset, + picture: o.picture, + }); + } + + /// Payload bytes of one chunk (`pos`: where the reader restarts to find them). + pub fn feed(&mut self, pos: u64, data: &[u8]) { + if data.is_empty() { + return; + } + // PES-delimited modes resolve a unit's position from `pos` right here and never look a + // chunk up again, so remembering chunks there would only grow without bound (NAL streams + // keep theirs: an access unit can start inside a PES; `keep_from` prunes them) + if matches!(self.mode, Mode::MpegVideo | Mode::Audio(_) | Mode::Nal { .. }) { + self.chunks.push_back((pos, self.es_pos)); + } + if let Some(o) = self.open.as_mut() + && o.pos == u64::MAX + && o.start == self.es_pos + { + o.pos = pos; + o.offset = 0; + } + let start = self.es_pos; + self.es_pos += data.len() as u64; + match self.mode { + Mode::Pes => {} + Mode::Nal { hevc } => self.scan_nal(start, data, hevc), + Mode::MpegVideo => self.scan_mpeg(start, data), + Mode::Audio(sync) => self.scan_audio(start, data, sync), + } + // forget chunks no pending position can refer to + let keep_from = self.buf_start.min(self.pending.map_or(u64::MAX, |p| p.unit_start.min(p.at))).min(self.header_start.unwrap_or(u64::MAX)); + while self.chunks.len() > 2 && self.chunks[1].1 <= keep_from { + self.chunks.pop_front(); + } + } + + /// Append to the window and return it with its start. + fn window(&mut self, start: u64, data: &[u8]) { + if self.buf.is_empty() || self.buf_start + self.buf.len() as u64 != start { + self.buf.clear(); + self.buf_start = start; + } + self.buf.extend_from_slice(data); + } + + /// Keep only the last `keep` bytes of the window. + fn trim(&mut self, keep: usize) { + if self.buf.len() > keep { + let drop = self.buf.len() - keep; + self.buf.drain(..drop); + self.buf_start += drop as u64; + } + } + + /// Start codes in the window from index `from`: (window index of `00 00 01`, code). + fn next_start_code(&self, from: usize) -> Option { + let b = &self.buf; + let mut i = from; + while i + 3 <= b.len() { + match b[i + 2] { + 1 if b[i] == 0 && b[i + 1] == 0 => return Some(i), + 0 => i += 1, + _ => i += 3, + } + } + None + } + + fn scan_nal(&mut self, start: u64, data: &[u8], hevc: bool) { + self.window(start, data); + let mut i = 0; + let mut blocked = None; + while let Some(k) = self.next_start_code(i) { + // NAL header byte(s) and the first byte of the slice header needed + if k + if hevc { 6 } else { 5 } > self.buf.len() { + blocked = Some(k); + break; + } + // a four-byte start code begins with the zero_byte that belongs to its NAL + let at = self.buf_start + k as u64 - (k > 0 && self.buf.get(k - 1) == Some(&0)) as u64; + let (Some(&h), Some(&h1)) = (self.buf.get(k + 3), self.buf.get(k + if hevc { 5 } else { 4 })) else { + blocked = Some(k); + break; + }; + let (key, vcl, non_ref, au_start) = if hevc { + let t = (h >> 1) & 0x3F; + ( + (16..=23).contains(&t), + t < 32, + // sub-layer non-reference pictures (TRAIL_N, TSA_N, STSA_N, RADL_N, RASL_N) + t < 16 && t.is_multiple_of(2), + // first_slice_segment_in_pic_flag, or a NAL that precedes the first slice + // (VPS, SPS, PPS, AUD, prefix SEI, 41..=44, 48..=55; H.265 7.4.2.4.4) + if t < 32 { h1 & 0x80 != 0 } else { matches!(t, 32..=35 | 39 | 41..=44 | 48..=55) }, + ) + } else { + let t = h & 0x1F; + ( + t == 5, + (1..=5).contains(&t), + h & 0x60 == 0, + // first_mb_in_slice == 0 (ue(v) "1"), or a NAL that precedes the first slice + // (SEI, SPS, PPS, AUD, 14..=18; H.264 7.4.1.2.3) + match t { + 1 | 2 | 5 => h1 & 0x80 != 0, + 6..=9 | 14..=18 => true, + _ => false, + }, + ) + }; + let mut recovery = false; + if !hevc && h & 0x1F == 6 { + // the whole SEI NAL is needed to find a recovery point message among others + let end = self.next_start_code(k + 3); + if end.is_none() && self.buf.len() - k < SEI_WINDOW { + blocked = Some(k); + break; + } + let nal = self.buf.get(k + 3..end.unwrap_or(self.buf.len())).unwrap_or(&[]); + recovery = sei_has_recovery_point(nal); + } + // a unit holds one access unit: a new one starts at the first NAL that cannot belong + // to the slices already seen (a unit opened later than this NAL is left alone) + if au_start && self.open.as_ref().is_some_and(|o| o.vcl && at > o.start) { + self.begin_access_unit(at); + } + if let Some(o) = self.open.as_mut() + && at >= o.start + { + // a recovery point is a random-access point (the pictures before its recovery + // frame count completes may be imperfect) + o.key |= recovery; + if vcl { + o.key |= key; + // disposable while every slice is non-reference + o.disposable = if o.vcl { o.disposable && non_ref } else { non_ref }; + o.vcl = true; + } + } + i = k + 3; + } + self.keep_tail(i, blocked); + } + + /// Close the open unit and start the next access unit at `at`: it has no timestamp (the PES + /// timestamps belong to the first access unit of the packet). + fn begin_access_unit(&mut self, at: u64) { + self.close(at); + let (pos, offset) = self.locate(at).unwrap_or((u64::MAX, 0)); + self.open = Some(Open { start: at, pos, offset, pts: None, dts: None, key: false, disposable: false, vcl: false, picture: None }); + } + + /// Drop the scanned part of the window: keep from a start code waiting for more bytes, or + /// the last bytes a split start code may begin in. + fn keep_tail(&mut self, i: usize, blocked: Option) { + let len = self.buf.len(); + let resume = blocked.unwrap_or_else(|| i.max(len.saturating_sub(3))); + self.trim(len - resume.min(len)); + } + + fn scan_mpeg(&mut self, start: u64, data: &[u8]) { + self.window(start, data); + let mut i = 0; + let mut blocked = None; + while let Some(k) = self.next_start_code(i) { + // up to 5 bytes after the code are inspected + if k + 9 > self.buf.len() { + blocked = Some(k); + break; + } + let at = self.buf_start + k as u64; + let code = self.buf[k + 3]; + let p = [self.buf[k + 4], self.buf[k + 5], self.buf[k + 6], self.buf[k + 7], self.buf[k + 8]]; + match code { + 0xB3 => { + self.settle_pending(); + self.header_start.get_or_insert(at); + self.seq_in_unit = true; + } + 0xB8 => { + self.settle_pending(); + self.header_start.get_or_insert(at); + self.gop = Some((p[3] & 0x40 != 0, p[3] & 0x20 != 0)); + } + 0x00 => { + self.settle_pending(); + let info = PictureInfo { + coding_type: (p[1] >> 3) & 7, + temporal_reference: ((p[0] as u16) << 2) | (p[1] >> 6) as u16, + structure: 3, + progressive_frame: true, + sequence_header: self.seq_in_unit, + gop: self.gop.is_some(), + closed_gop: self.gop.is_some_and(|g| g.0), + broken_link: self.gop.is_some_and(|g| g.1), + ..Default::default() + }; + let unit_start = self.header_start.take().unwrap_or(at); + self.seq_in_unit = false; + self.gop = None; + self.pending = Some(PendingPicture { unit_start, at, info }); + } + 0xB5 if p[0] >> 4 == 8 => { + if let Some(pp) = self.pending.as_mut() { + pp.info.structure = p[2] & 3; + pp.info.top_field_first = p[3] & 0x80 != 0; + pp.info.repeat_first_field = p[3] & 0x02 != 0; + pp.info.progressive_frame = p[4] & 0x80 != 0; + } + self.settle_pending(); + } + 0xB5 | 0xB2 => {} + _ => self.settle_pending(), + } + i = k + 3; + } + self.keep_tail(i, blocked); + } + + /// A picture's structure is known: start a unit, or join a first field. + fn settle_pending(&mut self) { + let Some(pp) = self.pending.take() else { return }; + let s = pp.info.structure; + if let Some(first) = self.open_field + && s != 3 + && first != s + && s != 0 + { + self.open_field = None; + return; + } + self.close(pp.unit_start); + let (pos, offset) = self.locate(pp.unit_start).unwrap_or((u64::MAX, 0)); + let mut o = Open { + start: pp.unit_start, + pos, + offset, + pts: None, + dts: None, + key: matches!(pp.info.coding_type, 1 | 4), + disposable: pp.info.coding_type == 3, + vcl: true, + picture: Some(pp.info), + }; + if let Some((pts, dts)) = self.timestamps_at(pp.at) { + o.pts = pts; + o.dts = dts; + } + self.open = Some(o); + self.open_field = (s == 1 || s == 2).then_some(s); + } + + fn scan_audio(&mut self, start: u64, data: &[u8], sync: AudioSync) { + let end = start + data.len() as u64; + // skip whole chunks inside the current frame + if let Some(n) = self.next_frame + && n >= end + { + return; + } + self.window(start, data); + loop { + let at = match self.next_frame { + Some(n) if n >= self.buf_start => n, + _ => self.buf_start, + }; + let i = (at - self.buf_start) as usize; + if i + 8 > self.buf.len() { + // not enough bytes for a header yet + self.trim(self.buf.len().saturating_sub(i)); + return; + } + match frame_len(sync, &self.buf[i..]) { + Some(len) if self.next_frame == Some(at) || self.confirmed(sync, i, len) => { + // an E-AC-3 dependent substream or further program joins the access unit of + // the independent substream 0 before it, as in MP4 and Matroska + if !(sync == AudioSync::Ac3 && self.open.is_some() && eac3_joins_unit(&self.buf[i..])) { + self.close(at); + let (pos, offset) = self.locate(at).unwrap_or((u64::MAX, 0)); + let mut o = Open { start: at, pos, offset, pts: None, dts: None, key: true, disposable: false, vcl: false, picture: None }; + if let Some((pts, dts)) = self.timestamps_at(at) { + o.pts = pts; + o.dts = dts; + } + self.open = Some(o); + } + self.next_frame = Some(at + len as u64); + if at + len as u64 >= self.buf_start + self.buf.len() as u64 { + self.buf.clear(); + self.buf_start = at + len as u64; + return; + } + } + _ => { + // lost sync: look for the next header from the following byte + self.next_frame = None; + self.skipped += 1; + let next = self.buf[i + 1..].iter().position(|&b| { + matches!((sync, b), (AudioSync::Ac3, 0x0B) | (AudioSync::Loas, 0x56) | (AudioSync::Mpa | AudioSync::Adts, 0xFF)) + }); + match next { + Some(d) => { + self.skipped += d as u64; + // Keep a cursor into the current window rather than draining its + // prefix for every false sync candidate. + self.next_frame = Some(at + 1 + d as u64); + } + None => { + self.skipped += (self.buf.len() - i - 1) as u64; + // A split header can begin in the final few bytes; compact once. + self.trim(7); + return; + } + } + } + } + } + } + + /// A header found while hunting for sync is accepted when the next frame's header follows + /// it (or the data ends before it). + fn confirmed(&self, sync: AudioSync, i: usize, len: usize) -> bool { + let j = i + len; + if j + 8 > self.buf.len() { + return true; + } + frame_len(sync, &self.buf[j..]).is_some() + } + + /// End of stream: close the open unit. + pub fn finish(&mut self) { + self.settle_pending(); + let end = match (self.mode, self.next_frame) { + (Mode::Audio(_), Some(n)) => n.min(self.es_pos), + _ => self.es_pos, + }; + if let (Mode::Audio(_), Some(o)) = (self.mode, self.open) { + // a truncated last frame is dropped + if self.next_frame.is_some_and(|n| n > self.es_pos) { + self.skipped += self.es_pos - o.start; + self.open = None; + } + } + self.close(end); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn audio_frame_sizes() { + // MPEG-1 layer II 48 kHz 192 kb/s: 576 bytes; layer III 44.1 kHz 128 kb/s with padding: 418 + assert_eq!(mpa_header(&[0xFF, 0xFD, 0xA4, 0x00]), Some((576, 48_000, 1152))); + assert_eq!(mpa_header(&[0xFF, 0xFB, 0x92, 0x00]), Some((418, 44_100, 1152))); + // MPEG-2 layer III 24 kHz 64 kb/s: 192 bytes, 576 samples + assert_eq!(mpa_header(&[0xFF, 0xF3, 0x84, 0x00]), Some((192, 24_000, 576))); + // AC-3 448 kb/s at 48 kHz = 1792 bytes; 44.1 kHz 192 kb/s odd code = 836 bytes + assert_eq!(ac3_frame_bytes(0, 30), Some(1792)); + assert_eq!(ac3_frame_bytes(1, 21), Some(836)); + assert_eq!(ac3_frame_bytes(2, 0), Some(192)); + } + + #[test] + fn eac3_dependent_substreams_join_the_access_unit() { + // 16-byte E-AC-3 syncframes (frmsiz 7), 48 kHz six blocks stereo: independent, dependent + // (7.1 extension), independent, dependent + let frame = |strmtyp: u8| { + let mut f = vec![0x0B, 0x77, strmtyp << 6, 7, (3 << 4) | (2 << 1), 16 << 3]; + f.resize(16, 0); + f + }; + let es: Vec = [frame(0), frame(1), frame(0), frame(1)].concat(); + let mut s = Splitter::new(&Codec::Eac3); + s.feed(0, &es); + s.finish(); + assert_eq!(s.units.iter().map(|u| u.size).collect::>(), vec![32, 32]); + assert_eq!(frame_info(&Codec::Eac3, &es).map(|f| f.samples), Some(1536)); + // AC-3 frames are never joined + assert!(!eac3_joins_unit(&[0x0B, 0x77, 0x40, 0, 0, 8 << 3])); + } + + #[test] + fn ac3_header_prefix_of_six_bytes_is_not_indexed_out_of_bounds() { + // exactly the six bytes frame_len() guarantees: byte 6 (acmod) is missing + assert_eq!(frame_info(&Codec::Ac3, &[0x0B, 0x77, 0, 0, 0x1E, 0x40]), None); + assert_eq!(frame_info(&Codec::Ac3, &[0x0B, 0x77, 0, 0, 0x1E, 0x40, 0]), None); + } + + #[test] + fn eac3_reserved_sample_rate_is_refused() { + // fscod 3 and fscod2 3: no sample rate + let mut f = vec![0x0B, 0x77, 0, 7, 0xF0, 16 << 3]; + f.resize(16, 0); + assert_eq!(frame_info(&Codec::Eac3, &f), None); + f[4] = 0xE0; // fscod2 2: 16 kHz + assert_eq!(frame_info(&Codec::Eac3, &f).map(|i| i.sample_rate), Some(16_000)); + } + + #[test] + fn four_byte_eac3_frames_make_no_units() { + let es: Vec = [0x0B, 0x77, 0x00, 0x01].repeat(10_000); + let mut s = Splitter::new(&Codec::Eac3); + s.feed(0, &es); + s.finish(); + assert!(s.units.is_empty(), "{} units from 40000 noise bytes", s.units.len()); + } + + #[test] + fn pes_delimited_modes_do_not_retain_chunks() { + for codec in [Codec::LpcmBluray, Codec::Unknown(0x99)] { + let mut s = Splitter::new(&codec); + for k in 0..5000u64 { + s.pes_start(Some(k as i64), None); + s.feed(k * 188, &[0x11; 184]); + } + s.finish(); + assert!(s.chunks.len() <= 2, "{codec:?}: {} chunks retained", s.chunks.len()); + assert_eq!(s.units.len(), 5000, "{codec:?}"); + assert!(s.units.iter().enumerate().all(|(k, u)| u.pos == k as u64 * 188 && u.offset == 0), "{codec:?}"); + } + } + + #[test] + fn unit_table_is_capped() { + let frame = |k: u8| { + let mut v = vec![0x0B, 0x77, 0, 7, (3 << 4) | (2 << 1), 16 << 3]; + v.resize(16, k); + v + }; + let mut s = Splitter::new(&Codec::Eac3); + s.max_units = 10; + let es: Vec = (0..100).flat_map(frame).collect(); + s.feed(0, &es); + s.finish(); + assert_eq!(s.units.len(), 10); + assert!(s.truncated); + } +} diff --git a/crates/media/src/mpegts/hostile_tests.rs b/crates/media/src/mpegts/hostile_tests.rs new file mode 100644 index 000000000..3990acaef --- /dev/null +++ b/crates/media/src/mpegts/hostile_tests.rs @@ -0,0 +1,209 @@ +//! Hostile input: streams of every codec filled with start codes and sync words at random, +//! damaged containers, and the public frame-header helpers on arbitrary bytes. Errors and empty +//! tables are fine; a panic (an index out of range, an overflow in a debug build) is not. + +use super::tests::{TsMux, pes}; +use super::*; + +/// Deterministic xorshift. +struct Rng(u64); + +impl Rng { + fn next(&mut self) -> u64 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + self.0 + } + fn below(&mut self, n: usize) -> usize { + (self.next() % n.max(1) as u64) as usize + } +} + +/// Random bytes with the patterns the splitters look for scattered through them: start codes of +/// every kind, MPEG audio / ADTS / LOAS / AC-3 sync words, with header bytes at the extremes. +/// At least one byte, at most about `max`. +fn soup(r: &mut Rng, max: usize) -> Vec { + let n = 1 + r.below(max); + let mut v = Vec::with_capacity(n + 16); + while v.len() < n { + let edge = [0x00u8, 0xFF, 0x80, 0x7F, r.next() as u8][r.below(5)]; + match r.below(12) { + 0 => v.extend([0, 0, 1, r.next() as u8]), + 1 => v.extend([0, 0, 1, [0x00, 0xB3, 0xB5, 0xB8, 0x09, 0x65, 0x41, 0x06, 0x26, 0x40, 0x46][r.below(11)], edge, edge]), + 2 => v.extend([0, 0, 0, 1, r.next() as u8, edge]), + 3 => v.extend([0xFF, 0xF0 | (r.next() as u8 & 0x0F), r.next() as u8, edge, edge, edge, edge]), + 4 => v.extend([0x0B, 0x77, edge, edge, r.next() as u8, r.next() as u8, edge]), + 5 => v.extend([0x56, 0xE0 | (r.next() as u8 & 0x1F), r.next() as u8]), + 6 => v.extend([edge; 3]), + _ => v.push(r.next() as u8), + } + } + v +} + +/// Opens and reads; returns how many units the tables hold. +fn open_and_read_all(data: &[u8]) -> usize { + let Ok(f) = open(data) else { return 0 }; + let _ = (f.format.name(), f.pcr_range); + for (s, st) in f.streams.iter().enumerate() { + let _ = (st.codec_params(), st.codec.name(), st.pts_range()); + let n = st.units.len(); + // every unit of a small table, a spread of a large one + for i in (0..n).step_by((n / 64).max(1)).chain(n.checked_sub(1)) { + let _ = f.read_unit(data, s, i); + } + } + let _ = f.read_unit(data, f.streams.len(), 0); + f.streams.iter().map(|s| s.units.len()).sum() +} + +/// Every stream type the PMT maps, each on its own PID. +const TYPES: [u8; 16] = [0x01, 0x02, 0x03, 0x04, 0x0F, 0x11, 0x1B, 0x24, 0x80, 0x81, 0x83, 0x84, 0x87, 0x90, 0x06, 0x42]; + +fn soup_ts(r: &mut Rng, m2ts: bool, packets: usize) -> Vec { + let mut mux = TsMux::new(m2ts); + let streams: Vec<(u8, u16)> = TYPES.iter().enumerate().map(|(i, &t)| (t, 0x200 + i as u16)).collect(); + mux.psi(&streams); + for _ in 0..packets { + let (_, pid) = streams[r.below(streams.len())]; + let payload = soup(r, 700); + let pts = (r.below(3) > 0).then(|| r.next() as i64 & ((1 << 33) - 1)); + let dts = pts.filter(|_| r.below(2) == 0).map(|p| p.wrapping_sub(3600) & ((1 << 33) - 1)); + mux.pes(pid, &pes([0xE0, 0xC0, 0xBD, 0xFD][r.below(4)], pts, dts, &payload, r.below(2) == 0)); + } + mux.out +} + +fn soup_ps(r: &mut Rng, mpeg1: bool, packets: usize) -> Vec { + let mut out = Vec::new(); + for _ in 0..packets { + if mpeg1 { + out.extend([0, 0, 1, 0xBA, 0x21, 0, 1, 0, 1, 0x80, 0, 1]); + } else { + out.extend([0, 0, 1, 0xBA, 0x44, 0, 4, 0, 4, 1, 1, 0x89, 0xC3, 0xF8 | r.below(8) as u8]); + out.extend(std::iter::repeat_n(0xFF, r.below(8))); + } + let mut payload = soup(r, 1500); + let id = [0xE0, 0xE1, 0xC0, 0xC1, 0xBD, 0xBD, 0xBD, 0xBB, 0xBC, 0xBE, 0xBF][r.below(11)]; + if id == 0xBD { + // DVD sub-streams: AC-3, DTS, LPCM (every format byte), sub-pictures + payload[0] = [0x80, 0x88, 0xA0, 0xA1, 0x20, r.next() as u8][r.below(6)]; + } + let pts = (r.below(3) > 0).then(|| r.next() as i64 & ((1 << 33) - 1)); + out.extend(pes(id, pts, None, &payload, true)); + } + out.extend([0, 0, 1, 0xB9]); + out +} + +#[test] +fn streams_of_start_codes_and_sync_words_never_panic() { + let mut r = Rng(0x1234_5678_9ABC_DEF1); + let (mut ts, mut ps) = (0, 0); + for round in 0..40 { + ts += open_and_read_all(&soup_ts(&mut r, round % 2 == 1, 150)); + ps += open_and_read_all(&soup_ps(&mut r, round % 2 == 1, 60)); + } + // the streams reach the splitters: they are cut into units, not rejected up front + assert!(ts > 2000 && ps > 500, "{ts} transport stream units, {ps} program stream units"); +} + +#[test] +fn damaged_containers_never_panic() { + let mut r = Rng(0xDEAD_BEEF_0BAD_F00D); + let bases = [soup_ts(&mut r, false, 60), soup_ts(&mut r, true, 60), soup_ps(&mut r, false, 40), soup_ps(&mut r, true, 40)]; + for base in &bases { + for _ in 0..300 { + let mut d = base.clone(); + for _ in 0..1 + r.below(24) { + let at = r.below(d.len()); + // header bytes at their extremes, lengths that run past the packet, lost sync + d[at] = [0x00, 0xFF, 0x47, 0x01, 0x80, 0xBA, r.next() as u8][r.below(7)]; + } + match r.below(4) { + 0 => d.truncate(r.below(d.len())), + // a run cut out of the middle: the packet grid shifts + 1 => { + let at = r.below(d.len()); + d.drain(at..(at + r.below(300)).min(d.len())); + } + _ => {} + } + open_and_read_all(&d); + } + } +} + +/// Program-specific information with every length field wrong. +#[test] +fn damaged_psi_sections_never_panic() { + let mut r = Rng(0x0F0F_1234_5555_AAAA); + for _ in 0..3000 { + let mut mux = TsMux::new(false); + // sections the muxer signs with a valid CRC, so the parsers see the damage + let mut pat = vec![0x00, 0xB0, 13, 0, 1, 0xC1, 0, 0, 0, 1, 0xE1, 0x00]; + let mut pmt = vec![0x02, 0xB0, 0, 0, 1, 0xC1, 0, 0, 0xE1, 0x00, 0xF0, 0]; + for _ in 0..r.below(6) { + pmt.extend([TYPES[r.below(TYPES.len())], 0xE2, r.next() as u8, 0xF0 | r.below(16) as u8, r.next() as u8]); + pmt.extend(soup(&mut r, 12)); + } + let len = pmt.len() - 3 + 4; + pmt[1] = 0xB0 | (len >> 8) as u8; + pmt[2] = len as u8; + for sec in [&mut pat, &mut pmt] { + for _ in 0..r.below(3) { + let at = r.below(sec.len()); + sec[at] = [0x00, 0xFF, 0x0F, 0xF0, r.next() as u8][r.below(5)]; + } + } + mux.section(0, &pat); + mux.section(0x100, &pmt); + for i in 0..8u16 { + mux.pes(0x200 + i, &pes(0xE0, Some(0), None, &soup(&mut r, 200), false)); + } + // enough packets for the sniffer + open_and_read_all(&mux.out); + } +} + +#[test] +fn frame_header_helpers_never_panic_on_any_bytes() { + let mut r = Rng(0x5EED_5EED_5EED_5EED); + let codecs = [Codec::MpegAudio, Codec::AacAdts, Codec::AacLatm, Codec::Ac3, Codec::Eac3, Codec::H264, Codec::LpcmDvd]; + for _ in 0..20_000 { + let b = soup(&mut r, 24); + for cut in 0..=b.len() { + let b = &b[..cut]; + for c in &codecs { + if let Some(n) = frame_bytes(c, b) { + assert!(n > 0, "{c:?}: an empty frame would never advance"); + } + let _ = frame_info(c, b); + } + let _ = (sei_has_recovery_point(b), pes::parse(b), sniff(b)); + let _ = ts::codec_of(b.first().copied().unwrap_or(6), b, cut % 2 == 0); + } + } + // every MPEG audio, ADTS and AC-3 header value in its first four bytes + for h in 0..=0xFFFFu32 { + for sync in [[0xFFu8, 0xE0], [0x0B, 0x77]] { + let b = [sync[0], sync[1] | (h >> 12) as u8, (h >> 4) as u8, (h << 4) as u8, h as u8, (h >> 8) as u8, 0xFF, 0x00]; + for c in &codecs[..5] { + let _ = (frame_bytes(c, &b), frame_info(c, &b)); + } + } + } +} + +/// A pack header is looked for behind at most 64 zero bytes. +#[test] +fn a_pack_header_behind_too_many_zero_bytes_is_not_sniffed() { + let pack = [0u8, 0, 1, 0xBA, 0x44, 0, 4, 0, 4, 1, 1, 0x89, 0xC3, 0xF8]; + let behind = |zeros: usize| [vec![0u8; zeros], pack.to_vec()].concat(); + assert_eq!(sniff(&behind(0)), Some(Format::Ps { mpeg1: false })); + assert_eq!(sniff(&behind(62)), Some(Format::Ps { mpeg1: false })); + assert_eq!(sniff(&behind(63)), None); + assert_eq!(sniff(&behind(5000)), None); + assert!(matches!(open(&behind(5000)), Err(Error::NotMpeg))); +} diff --git a/crates/media/src/mpegts/mod.rs b/crates/media/src/mpegts/mod.rs new file mode 100644 index 000000000..2e100f336 --- /dev/null +++ b/crates/media/src/mpegts/mod.rs @@ -0,0 +1,458 @@ +//! Clean-room MPEG-2 Systems demuxer (ITU-T H.222.0 | ISO/IEC 13818-1). +//! +//! - **Transport Stream**: 188-byte packets, 192-byte BDAV / AVCHD packets (`.m2ts`, `.mts`: +//! a 4-byte arrival timestamp before each packet) and 204-byte packets (Reed-Solomon parity). +//! PAT / PMT (sections reassembled across packets, CRC checked), PES reassembly, PCR, PTS/DTS +//! (33-bit wrap-around unwrapped), discontinuity indicators. +//! - **Program Stream** (`.mpg`, `.vob`, `.mod`, `.tod`): MPEG-2 and ISO/IEC 11172-1 (MPEG-1) +//! packs, system headers, program stream maps, DVD private stream 1 sub-streams (AC-3, DTS, +//! LPCM). +//! - **Elementary streams** are split into access units on open — coded frames for MPEG-1/2 +//! video (picture start codes, field pairs), H.264 and HEVC (access unit delimiters / first +//! slice of a picture), audio frames for MPEG audio, ADTS / LATM AAC, AC-3 and E-AC-3, PES +//! packets for LPCM — each with its PTS/DTS (a PES timestamp belongs to the first access unit +//! that starts in the packet), random-access flag and, for MPEG video, picture type, temporal +//! reference and field structure. That index is the seeking table. +//! - Units are read back on demand by re-reading the file from the unit's first packet, so the +//! index costs a few dozen bytes per access unit. +//! +//! The module does no decoding. No dependencies beyond `std`, no `unsafe`, builds for +//! `wasm32-unknown-unknown`. +//! +//! From FilmCraft's `filmcraft-mpegts` (storytold/filmcraft@4bec22b). Public items keep their names. +//! `ByteSource` is the crate's one source trait. + +#![warn(missing_docs)] + +mod es; +pub use es::{FrameInfo, frame_bytes, frame_info, sei_has_recovery_point}; +mod pes; +mod ps; +mod ts; + +use std::fmt; +use std::io; + +pub use crate::source::{ByteSource, ReadSeekSource}; + +/// Errors from reading a transport or program stream. +#[non_exhaustive] +#[derive(Debug)] +pub enum Error { + /// Neither a transport stream nor a program stream. + NotMpeg, + /// Reading the source failed. + Io(io::Error), + /// Structurally invalid data. + Invalid(String), +} + +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Error::NotMpeg => write!(f, "not an MPEG transport or program stream"), + Error::Io(e) => write!(f, "I/O: {e}"), + Error::Invalid(s) => write!(f, "invalid stream: {s}"), + } + } +} + +impl std::error::Error for Error { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Error::Io(e) => Some(e), + _ => None, + } + } +} + +impl From for Error { + fn from(e: io::Error) -> Self { + Error::Io(e) + } +} + +/// Result alias for this module. +pub type Result = std::result::Result; + +/// Read as much as is available at `offset` (up to `buf.len()`); returns the count. +pub(crate) fn read_some(src: &S, offset: u64, buf: &mut [u8]) -> io::Result { + let len = src.len(); + if offset >= len { + return Ok(0); + } + let n = ((len - offset) as usize).min(buf.len()); + src.read_at(offset, &mut buf[..n])?; + Ok(n) +} + +/// The container layout. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Format { + /// Transport stream with `packet_size`-byte packets whose 0x47 sync byte is at `sync_offset` + /// (4 for 192-byte BDAV packets). + Ts { + /// 188, 192 or 204. + packet_size: usize, + /// Offset of the sync byte in a packet. + sync_offset: usize, + }, + /// Program stream (`mpeg1`: ISO/IEC 11172-1 packs). + Ps { + /// ISO/IEC 11172-1 packs. + mpeg1: bool, + }, +} + +impl Format { + /// A display name: `"MPEG-2 TS"`, `"MPEG-2 PS"`… + pub fn name(&self) -> &'static str { + match self { + Format::Ts { packet_size: 192, .. } => "MPEG-2 TS (BDAV/AVCHD)", + Format::Ts { .. } => "MPEG-2 TS", + Format::Ps { mpeg1: true } => "MPEG-1 System", + Format::Ps { mpeg1: false } => "MPEG-2 PS", + } + } +} + +/// Recognise a transport stream (five consecutive sync bytes at one of the packet sizes, within +/// the first packets) or a program stream (a pack header). +pub fn sniff(head: &[u8]) -> Option { + for (size, off) in [(188usize, 0usize), (192, 4), (204, 0)] { + for start in 0..size.min(head.len()) { + let ok = (0..5).all(|k| head.get(start + off + k * size) == Some(&0x47)); + if ok { + return Some(Format::Ts { packet_size: size, sync_offset: off }); + } + } + } + // a pack header, possibly after a few zero bytes + let z = head.iter().take(64).take_while(|&&b| b == 0).count(); + let p = z.saturating_sub(2); + if head.len() >= p + 5 && head[p..p + 4] == [0, 0, 1, 0xBA] { + return Some(Format::Ps { mpeg1: head[p + 4] >> 6 != 1 }); + } + None +} + +/// What kind of media a stream carries. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Kind { + /// Pictures. + Video, + /// Sound. + Audio, + /// Subtitles (PGS, DVB, teletext). + Subtitle, + /// Anything else. + Other, +} + +/// What an elementary stream carries. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Codec { + /// MPEG-1 video (ISO/IEC 11172-2). + Mpeg1Video, + /// MPEG-2 video (H.262). + Mpeg2Video, + /// H.264 / AVC, Annex B. + H264, + /// H.265 / HEVC, Annex B. + Hevc, + /// MPEG-1/2 audio (layers I-III; the layer is in each frame header). + MpegAudio, + /// AAC in ADTS frames. + AacAdts, + /// AAC in LATM/LOAS frames. + AacLatm, + /// AC-3. + Ac3, + /// E-AC-3. + Eac3, + /// DTS. + Dts, + /// Dolby TrueHD. + TrueHd, + /// Blu-ray / AVCHD LPCM (a 4-byte header per PES packet, big-endian samples). + LpcmBluray, + /// DVD LPCM (program stream private stream 1; format in [`Stream::lpcm`]). + LpcmDvd, + /// Subtitles; the text names the format. + Subtitle(&'static str), + /// Anything else: the stream type. + Unknown(u8), +} + +impl Codec { + /// A display name: `"H.264"`, `"AAC (ADTS)"`, `"stream type 0x42"`. + pub fn name(&self) -> String { + match self { + Codec::Mpeg1Video => "MPEG-1 Video".into(), + Codec::Mpeg2Video => "MPEG-2 Video".into(), + Codec::H264 => "H.264".into(), + Codec::Hevc => "HEVC".into(), + Codec::MpegAudio => "MPEG Audio".into(), + Codec::AacAdts => "AAC (ADTS)".into(), + Codec::AacLatm => "AAC (LATM)".into(), + Codec::Ac3 => "AC-3".into(), + Codec::Eac3 => "E-AC-3".into(), + Codec::Dts => "DTS".into(), + Codec::TrueHd => "Dolby TrueHD".into(), + Codec::LpcmBluray => "LPCM (Blu-ray)".into(), + Codec::LpcmDvd => "LPCM (DVD)".into(), + Codec::Subtitle(s) => (*s).into(), + Codec::Unknown(t) => format!("stream type 0x{t:02x}"), + } + } + /// The crate-wide identifier. MPEG-1 video reports [`CodecId::Mpeg2Video`] (an MPEG-2 + /// decoder decodes it), MPEG audio reports [`CodecId::Mp2`] whatever its layer (the layer is + /// in each frame header: [`frame_info`]'s `variant`), and DTS, TrueHD and subtitles have no + /// identifier. + /// + /// [`CodecId::Mpeg2Video`]: crate::CodecId::Mpeg2Video + /// [`CodecId::Mp2`]: crate::CodecId::Mp2 + pub fn codec_id(&self) -> crate::CodecId { + use crate::CodecId as Id; + match self { + Codec::Mpeg1Video | Codec::Mpeg2Video => Id::Mpeg2Video, + Codec::H264 => Id::H264, + Codec::Hevc => Id::Hevc, + Codec::MpegAudio => Id::Mp2, + Codec::AacAdts | Codec::AacLatm => Id::Aac, + Codec::Ac3 => Id::Ac3, + Codec::Eac3 => Id::Eac3, + Codec::LpcmBluray | Codec::LpcmDvd => Id::Pcm, + Codec::Dts | Codec::TrueHd | Codec::Subtitle(_) | Codec::Unknown(_) => Id::Unknown, + } + } + + /// The kind of media. + pub fn kind(&self) -> Kind { + match self { + Codec::Mpeg1Video | Codec::Mpeg2Video | Codec::H264 | Codec::Hevc => Kind::Video, + Codec::MpegAudio + | Codec::AacAdts + | Codec::AacLatm + | Codec::Ac3 + | Codec::Eac3 + | Codec::Dts + | Codec::TrueHd + | Codec::LpcmBluray + | Codec::LpcmDvd => Kind::Audio, + Codec::Subtitle(_) => Kind::Subtitle, + Codec::Unknown(_) => Kind::Other, + } + } +} + +/// MPEG-1/2 video picture properties of an access unit (its first picture). +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct PictureInfo { + /// 1 I, 2 P, 3 B, 4 D. + pub coding_type: u8, + /// `temporal_reference`: display order within the GOP. + pub temporal_reference: u16, + /// 1 top field, 2 bottom field, 3 frame (two field pictures form one unit). + pub structure: u8, + /// `top_field_first` of the picture coding extension. + pub top_field_first: bool, + /// `repeat_first_field` (pulldown). + pub repeat_first_field: bool, + /// `progressive_frame`. + pub progressive_frame: bool, + /// A sequence header precedes the picture in this unit. + pub sequence_header: bool, + /// A GOP header precedes the picture (and its flags). + pub gop: bool, + /// `closed_gop` of that GOP header. + pub closed_gop: bool, + /// `broken_link` of that GOP header. + pub broken_link: bool, +} + +/// One access unit (a coded frame, an audio frame, or an LPCM PES packet). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Unit { + /// Presentation / decoding time stamps, 90 kHz, 33-bit wrap-around unwrapped. + pub pts: Option, + /// See `pts`. + pub dts: Option, + /// Decoding can start here (I / IDR / IRAP pictures; every audio frame). + pub key: bool, + /// No other unit depends on it (B pictures of MPEG-1/2 video, non-reference H.264/HEVC). + pub disposable: bool, + /// Bytes of elementary stream. + pub size: u32, + /// TS: byte offset of the packet holding the first byte; PS: byte offset of its PES packet. + pub(crate) pos: u64, + /// Offset of the first byte within that packet's elementary-stream payload. + pub(crate) offset: u32, + /// MPEG-1/2 video: the picture's properties. + pub picture: Option, +} + +/// DVD LPCM format (from the private stream 1 header). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct LpcmFormat { + /// Samples per second. + pub sample_rate: u32, + /// Channel count. + pub channels: u32, + /// Bits per sample (16, 20 or 24). + pub bits: u32, +} + +/// Where a stream's PES packets come from. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum StreamId { + /// Transport stream PID. + Pid(u16), + /// Program stream stream_id, and the private stream 1 sub-stream id. + Ps { + /// PES stream_id. + stream_id: u8, + /// Private stream 1 sub-stream id. + sub_id: Option, + }, +} + +/// One elementary stream and its access-unit table. +#[derive(Clone, Debug)] +pub struct Stream { + /// Where its PES packets come from. + pub id: StreamId, + /// PMT stream_type (TS), or the equivalent for program streams. + pub stream_type: u8, + /// What it carries. + pub codec: Codec, + /// ISO 639 language code from the PMT. + pub language: Option, + /// The access units, in stream order. + pub units: Vec, + /// DVD LPCM: the sample format. + pub lpcm: Option, + /// PES packets seen. + pub pes_packets: u64, + /// Bytes of payload ignored because the stream could not be split (e.g. before the first + /// frame header). + pub skipped_bytes: u64, +} + +impl Stream { + /// The kind of media. + pub fn kind(&self) -> Kind { + self.codec.kind() + } + /// What a decoder needs for this stream, in the container-neutral form the codec crates take: + /// the codec ([`Codec::codec_id`]) and, for DVD LPCM, the sample format (big-endian). These + /// streams carry their configuration in band, so `extradata` is empty: H.264 and HEVC are + /// Annex B with their parameter sets, AAC is ADTS or LATM frames, and a Blu-ray LPCM packet + /// starts with its own 4-byte format header. + pub fn codec_params(&self) -> crate::CodecParams { + let mut p = crate::CodecParams::new(self.codec.codec_id()); + if let Some(f) = self.lpcm { + let n = |v: u32| u16::try_from(v).unwrap_or(0); + p = p.with_sample_rate(f.sample_rate).with_channels(n(f.channels)).with_pcm(crate::PcmFormat::int(n(f.bits), true, true)); + } + p + } + /// First and last presentation timestamps (90 kHz) of the stream's units. + pub fn pts_range(&self) -> Option<(i64, i64)> { + let mut it = self.units.iter().filter_map(|u| u.pts); + let first = it.next()?; + Some(it.fold((first, first), |(a, b), p| (a.min(p), b.max(p)))) + } +} + +/// A demuxed file: its streams with their access-unit tables. +#[derive(Clone, Debug)] +pub struct File { + /// The container layout. + pub format: Format, + /// The elementary streams (TS: in PMT order). + pub streams: Vec, + /// Program number of the program the streams belong to (TS). + pub program: Option, + /// First and last PCR (27 MHz, unwrapped) or SCR. + pub pcr_range: Option<(i64, i64)>, + /// Problems found while scanning (lost sync, CRC errors, continuity errors…). + pub warnings: Vec, +} + +/// Largest unit `read_unit` returns. The biggest real units are intra pictures (a few MB at +/// 4K/8K broadcast rates); a stream without PES timestamps can make one unit of a whole PID. +pub const MAX_UNIT_BYTES: u32 = 128 << 20; + +/// Bytes reserved up front by `read_unit`; the buffer grows beyond this only as data is read. +const UNIT_PREALLOC: usize = 1 << 20; + +impl File { + /// The first stream of `kind` whose codec satisfies `f`. + pub fn find(&self, kind: Kind, f: impl Fn(&Codec) -> bool) -> Option { + self.streams.iter().position(|s| s.kind() == kind && !s.units.is_empty() && f(&s.codec)) + } + + /// Read access unit `i` of stream `s`. + pub fn read_unit(&self, src: &S, s: usize, i: usize) -> Result> { + let st = self.streams.get(s).ok_or_else(|| Error::Invalid("no such stream".into()))?; + let u = st.units.get(i).ok_or_else(|| Error::Invalid("no such unit".into()))?; + if u.size > MAX_UNIT_BYTES { + return Err(Error::Invalid(format!("unit {i} is {} bytes (limit {MAX_UNIT_BYTES})", u.size))); + } + // grow as bytes arrive: a damaged table must not reserve more than the file can supply + let mut out = Vec::with_capacity((u.size as usize).min(UNIT_PREALLOC)); + match self.format { + Format::Ts { packet_size, sync_offset } => ts::read(src, packet_size, sync_offset, st, u, &mut out)?, + Format::Ps { .. } => ps::read(src, st, u, &mut out)?, + } + if out.len() < u.size as usize { + return Err(Error::Invalid(format!("unit {i} truncated ({} of {} bytes)", out.len(), u.size))); + } + Ok(out) + } +} + +/// Scan a transport or program stream and index its access units. +pub fn open(src: &S) -> Result { + let mut head = vec![0u8; 64 * 1024]; + let n = read_some(src, 0, &mut head)?; + head.truncate(n); + match sniff(&head).ok_or(Error::NotMpeg)? { + Format::Ts { packet_size, sync_offset } => ts::scan(src, packet_size, sync_offset), + Format::Ps { mpeg1 } => ps::scan(src, mpeg1), + } +} + +/// Unwraps 33-bit timestamps into a monotonic-ish i64 timeline (relative to the first one seen). +#[derive(Default, Clone, Copy)] +pub(crate) struct Unwrap { + last: Option, + base: i64, +} + +impl Unwrap { + const WRAP: i64 = 1 << 33; + pub fn unwrap(&mut self, raw: i64) -> i64 { + let mut v = raw + self.base; + if let Some(l) = self.last { + if v < l - Self::WRAP / 2 { + self.base += Self::WRAP; + v += Self::WRAP; + } else if v > l + Self::WRAP / 2 && self.base > 0 { + // a timestamp from before the wrap (B pictures, slightly out-of-order audio) + v -= Self::WRAP; + return v; + } + } + self.last = Some(v); + v + } +} + +#[cfg(test)] +mod hostile_tests; +#[cfg(test)] +mod tests; +#[cfg(test)] +mod work_tests; diff --git a/crates/media/src/mpegts/pes.rs b/crates/media/src/mpegts/pes.rs new file mode 100644 index 000000000..06c7a2abf --- /dev/null +++ b/crates/media/src/mpegts/pes.rs @@ -0,0 +1,126 @@ +//! PES packet headers (H.222.0 §2.4.3.6; ISO/IEC 11172-1 §2.4.3.3 for MPEG-1 packets). + +/// A parsed PES header. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct PesHeader { + /// Bytes from the packet start code to the first payload byte. + pub header_len: usize, + /// Raw 33-bit timestamps. + pub pts: Option, + pub dts: Option, +} + +/// Stream ids whose packets have no optional header (§2.4.3.7). +pub(crate) fn plain_stream(id: u8) -> bool { + matches!(id, 0xBC | 0xBE | 0xBF | 0xF0 | 0xF1 | 0xF2 | 0xF8 | 0xFF) +} + +pub(crate) fn timestamp(b: &[u8]) -> i64 { + (((b[0] >> 1) & 7) as i64) << 30 | (b[1] as i64) << 22 | ((b[2] >> 1) as i64) << 15 | (b[3] as i64) << 7 | (b[4] >> 1) as i64 +} + +/// Parse a PES header at the start of `b` (which begins with `00 00 01 stream_id`). `None` when +/// `b` is too short to hold the whole header or it is malformed. +pub(crate) fn parse(b: &[u8]) -> Option { + if b.len() < 6 || b[0] != 0 || b[1] != 0 || b[2] != 1 { + return None; + } + if plain_stream(b[3]) { + return Some(PesHeader { header_len: 6, pts: None, dts: None }); + } + if b.len() >= 9 && b[6] & 0xC0 == 0x80 { + // ISO/IEC 13818-1 syntax + let flags = b[7]; + let header_len = 9 + b[8] as usize; + if b.len() < header_len { + return None; + } + let (mut pts, mut dts) = (None, None); + if flags & 0x80 != 0 && header_len >= 14 { + pts = Some(timestamp(&b[9..14])); + if flags & 0x40 != 0 && header_len >= 19 { + dts = Some(timestamp(&b[14..19])); + } + } + return Some(PesHeader { header_len, pts, dts }); + } + // ISO/IEC 11172-1 syntax: stuffing, STD buffer, timestamps + let mut p = 6; + let mut stuffing = 0; + while b.get(p) == Some(&0xFF) { + p += 1; + stuffing += 1; + if stuffing > 16 { + return None; + } + } + let c = *b.get(p)?; + if c >> 6 == 1 { + p += 2; + } + let c = *b.get(p)?; + let (mut pts, mut dts) = (None, None); + match c >> 4 { + 2 => { + pts = Some(timestamp(b.get(p..p + 5)?)); + p += 5; + } + 3 => { + pts = Some(timestamp(b.get(p..p + 5)?)); + dts = Some(timestamp(b.get(p + 5..p + 10)?)); + p += 10; + } + 0 if c == 0x0F => p += 1, + _ => return None, + } + Some(PesHeader { header_len: p, pts, dts }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ts_bytes(prefix: u8, v: i64) -> [u8; 5] { + [ + (prefix << 4) | (((v >> 30) & 7) as u8) << 1 | 1, + (v >> 22) as u8, + (((v >> 15) & 0x7F) as u8) << 1 | 1, + (v >> 7) as u8, + ((v & 0x7F) as u8) << 1 | 1, + ] + } + + #[test] + fn mpeg2_and_mpeg1_headers() { + let pts = 0x1_2345_6789i64; + let dts = 0x0_8765_4321i64; + let mut b = vec![0, 0, 1, 0xE0, 0, 0, 0x80, 0xC0, 10]; + b.extend(ts_bytes(3, pts)); + b.extend(ts_bytes(1, dts)); + b.push(0xAA); + let h = parse(&b).unwrap(); + assert_eq!((h.header_len, h.pts, h.dts), (19, Some(pts), Some(dts))); + // MPEG-1: stuffing, STD buffer, PTS only + let mut b = vec![0, 0, 1, 0xC0, 0, 20, 0xFF, 0xFF, 0x40, 0x20]; + b.extend(ts_bytes(2, 90_000)); + let h = parse(&b).unwrap(); + assert_eq!((h.header_len, h.pts, h.dts), (15, Some(90_000), None)); + assert_eq!(parse(&[0, 0, 1, 0xC0, 0, 1, 0x0F]).unwrap().header_len, 7); + assert_eq!(parse(&[0, 0, 1, 0xBE, 0, 4]).unwrap().header_len, 6); + assert!(parse(&[0, 0, 1, 0xE0, 0, 0, 0x80, 0x80]).is_none()); + } + + /// MPEG-1 packets: at most 16 stuffing bytes before the timestamps. + #[test] + fn mpeg1_stuffing_past_the_limit_is_refused() { + let with = |stuffing: usize| { + let mut b = vec![0, 0, 1, 0xC0, 0, 40]; + b.extend(std::iter::repeat_n(0xFF, stuffing)); + b.extend(ts_bytes(2, 90_000)); + parse(&b) + }; + assert_eq!(with(16).map(|h| (h.header_len, h.pts)), Some((6 + 16 + 5, Some(90_000)))); + assert_eq!(with(17), None); + assert_eq!(with(5000), None); + } +} diff --git a/crates/media/src/mpegts/ps.rs b/crates/media/src/mpegts/ps.rs new file mode 100644 index 000000000..41c5486a9 --- /dev/null +++ b/crates/media/src/mpegts/ps.rs @@ -0,0 +1,370 @@ +//! Program streams (H.222.0 §2.5; ISO/IEC 11172-1 §2.4.3): packs, system headers, the program +//! stream map, PES packets, and the DVD sub-streams of private stream 1. + +use std::collections::HashMap; + +use super::es::Splitter; +use super::pes; +use super::{ByteSource, Codec, Error, File, Format, LpcmFormat, Result, Stream, StreamId, Unit, Unwrap, read_some}; + +/// Sequential reader with a refillable window. +struct Win<'a, S: ByteSource + ?Sized> { + src: &'a S, + buf: Vec, + start: u64, + /// Bytes zero-filled by window refills (tests count them). + #[cfg(test)] + zero_filled: usize, +} + +impl<'a, S: ByteSource + ?Sized> Win<'a, S> { + fn new(src: &'a S) -> Self { + Self { + src, + buf: Vec::new(), + start: 0, + #[cfg(test)] + zero_filled: 0, + } + } + /// The bytes [pos, pos + n) (fewer at the end of the file). + fn get(&mut self, pos: u64, n: usize) -> Result<&[u8]> { + let have = pos >= self.start && pos + n as u64 <= self.start + self.buf.len() as u64; + if !have { + // never more than the file still holds past `pos`: `read_some` would not fill it, and + // zeroing a megabyte for every packet near the end of the file is wasted work + let remaining = self.src.len().saturating_sub(pos); + let want = (n.max(1 << 20) as u64).min(remaining) as usize; + #[cfg(test)] + { + self.zero_filled += want.saturating_sub(self.buf.len()); + } + self.buf.resize(want, 0); + let got = read_some(self.src, pos, &mut self.buf)?; + self.buf.truncate(got); + self.start = pos; + } + let i = (pos - self.start) as usize; + let end = (i + n).min(self.buf.len()); + Ok(&self.buf[i.min(end)..end]) + } +} + +/// How many bytes of a private stream 1 payload precede the elementary stream, by sub-stream id +/// (DVD-Video: AC-3 / DTS 4, LPCM 7, sub-pictures 1). +fn sub_header(sub: u8) -> usize { + match sub { + 0x80..=0x8F => 4, + 0xA0..=0xAF => 7, + _ => 1, + } +} + +fn codec_for(stream_id: u8, sub: Option, psm: &HashMap, mpeg1: bool) -> Codec { + if let Some(&t) = psm.get(&stream_id) { + let c = super::ts::codec_of(t, &[], false); + if !matches!(c, Codec::Unknown(_)) { + return c; + } + } + match (stream_id, sub) { + (0xE0..=0xEF, _) => { + if mpeg1 { + Codec::Mpeg1Video + } else { + Codec::Mpeg2Video + } + } + (0xC0..=0xDF, _) => Codec::MpegAudio, + (0xBD, Some(0x80..=0x87)) => Codec::Ac3, + (0xBD, Some(0x88..=0x8F)) => Codec::Dts, + (0xBD, Some(0xA0..=0xAF)) => Codec::LpcmDvd, + (0xBD, Some(0x20..=0x3F)) => Codec::Subtitle("DVD sub-pictures"), + (id, _) => Codec::Unknown(id), + } +} + +/// DVD LPCM header bytes after the sub-stream id (and the frame count / pointer). +fn lpcm_format(h: &[u8]) -> Option { + let b = *h.get(4)?; + let bits = [16, 20, 24, 0][(b >> 6) as usize]; + let sample_rate = [48_000, 96_000, 44_100, 32_000][((b >> 4) & 3) as usize]; + let channels = (b & 7) as u32 + 1; + (bits > 0).then_some(LpcmFormat { sample_rate, channels, bits }) +} + +struct EsState { + stream: Stream, + split: Splitter, +} + +/// Length of the packet (pack header, system header, PES…) starting at `b` (which begins with a +/// start code), `None` if it is not one. +fn packet_len(b: &[u8]) -> Option { + if b.len() < 6 || b[0] != 0 || b[1] != 0 || b[2] != 1 { + return None; + } + match b[3] { + 0xBA => { + if b[4] >> 6 == 1 { + let stuff = (*b.get(13)? & 7) as usize; + Some(14 + stuff) + } else if b[4] >> 4 == 2 { + // MPEG-1 pack header: start code, 5-byte SCR, 3-byte mux rate. A file cut inside + // it is not a pack header yet (`scan` reads the SCR from the first 9 bytes). + b.get(..12)?; + Some(12) + } else { + None + } + } + 0xB9 => Some(4), + 0xBB..=0xFF => Some(6 + u16::from_be_bytes([b[4], b[5]]) as usize), + _ => None, + } +} + +pub(crate) fn scan(src: &S, mpeg1: bool) -> Result { + let len = src.len(); + let mut win = Win::new(src); + let mut pos = 0u64; + let mut streams: HashMap<(u8, Option), EsState> = HashMap::new(); + let mut order = Vec::new(); + let mut psm: HashMap = HashMap::new(); + let mut unwrap = Unwrap::default(); + let mut scr: Option<(i64, i64)> = None; + let mut resyncs = 0u64; + while pos + 4 <= len { + let head = win.get(pos, 32)?; + let Some(plen) = packet_len(head) else { + // resynchronise on the next pack header or PES start code + resyncs += 1; + let mut p = pos + 1; + let found = loop { + let w = win.get(p, 1 << 16)?; + if w.len() < 4 { + break None; + } + if let Some(i) = w.windows(4).position(|x| x[0] == 0 && x[1] == 0 && x[2] == 1 && (x[3] == 0xBA || (0xBD..=0xEF).contains(&x[3]))) { + break Some(p + i as u64); + } + p += (w.len() - 3) as u64; + }; + match found { + Some(f) => { + pos = f; + continue; + } + None => break, + } + }; + let id = head[3]; + if id == 0xBA { + // SCR (MPEG-2: 33 bits base, 9 bits extension; MPEG-1: 33 bits) + let v = if head[4] >> 6 == 1 { + let b = &head[4..10]; + let base = (((b[0] >> 3) & 7) as i64) << 30 + | ((b[0] & 3) as i64) << 28 + | (b[1] as i64) << 20 + | ((b[2] >> 3) as i64) << 15 + | ((b[2] & 3) as i64) << 13 + | (b[3] as i64) << 5 + | (b[4] >> 3) as i64; + let ext = (((b[4] & 3) as i64) << 7) | (b[5] >> 1) as i64; + base * 300 + ext + } else { + pes::timestamp(&head[4..9]) * 300 + }; + scr = Some(scr.map_or((v, v), |(a, b): (i64, i64)| (a.min(v), b.max(v)))); + } else if id == 0xBC { + // program stream map: elementary stream id → stream_type + let body = win.get(pos, plen)?; + if body.len() >= 12 { + let info = u16::from_be_bytes([body[8], body[9]]) as usize; + let mut i = 10 + info; + if i + 2 <= body.len() { + let map_len = u16::from_be_bytes([body[i], body[i + 1]]) as usize; + i += 2; + let end = (i + map_len).min(body.len()); + while i + 4 <= end { + psm.insert(body[i + 1], body[i]); + i += 4 + u16::from_be_bytes([body[i + 2], body[i + 3]]) as usize; + } + } + } + } else if id == 0xBD || (0xC0..=0xEF).contains(&id) || id == 0xFD { + let pkt = win.get(pos, plen)?; + if pkt.len() < plen { + break; // truncated + } + if let Some(h) = pes::parse(pkt) { + let payload = &pkt[h.header_len.min(plen)..]; + let (sub, skip) = if id == 0xBD { + match payload.first() { + Some(&s) => (Some(s), sub_header(s).min(payload.len())), + None => (None, 0), + } + } else { + (None, 0) + }; + let key = (id, sub); + let e = streams.entry(key).or_insert_with(|| { + let codec = codec_for(id, sub, &psm, mpeg1); + let lpcm = if codec == Codec::LpcmDvd { lpcm_format(&payload[1.min(payload.len())..]) } else { None }; + order.push(key); + let stream = Stream { + id: StreamId::Ps { stream_id: id, sub_id: sub }, + stream_type: psm.get(&id).copied().unwrap_or(0), + codec: codec.clone(), + language: None, + units: Vec::new(), + lpcm, + pes_packets: 0, + skipped_bytes: 0, + }; + EsState { stream, split: Splitter::new(&codec) } + }); + e.stream.pes_packets += 1; + e.split.pes_start(h.pts.map(|t| unwrap.unwrap(t)), h.dts.map(|t| unwrap.unwrap(t))); + e.split.feed(pos, &payload[skip..]); + } + } + pos += plen as u64; + } + let mut warnings = Vec::new(); + if resyncs > 0 { + warnings.push(format!("resynchronised {resyncs} time(s)")); + } + let mut out = Vec::new(); + for key in order { + let Some(mut e) = streams.remove(&key) else { continue }; + e.split.finish(); + e.stream.units = std::mem::take(&mut e.split.units); + e.stream.skipped_bytes = e.split.skipped; + out.push(e.stream); + } + if out.is_empty() { + return Err(Error::Invalid("program stream without elementary streams".into())); + } + Ok(File { format: Format::Ps { mpeg1 }, streams: out, program: None, pcr_range: scr, warnings }) +} + +/// Collect unit `u` of stream `st`, starting at the PES packet at `u.pos`. +pub(crate) fn read(src: &S, st: &Stream, u: &Unit, out: &mut Vec) -> Result<()> { + let StreamId::Ps { stream_id, sub_id } = st.id else { return Err(Error::Invalid("not a PS stream".into())) }; + let need = u.size as usize; + let mut skip = u.offset as usize; + let mut pos = u.pos; + let len = src.len(); + let mut win = Win::new(src); + while out.len() < need && pos + 4 <= len { + let head = win.get(pos, 32)?; + let Some(plen) = packet_len(head) else { + // corrupt: stop at what we have + return Ok(()); + }; + if head[3] == stream_id { + let pkt = win.get(pos, plen)?; + if let Some(h) = pes::parse(pkt) { + let payload = &pkt[h.header_len.min(pkt.len())..]; + let matches = stream_id != 0xBD || payload.first().copied() == sub_id; + if matches { + let s0 = if stream_id == 0xBD { sub_id.map_or(0, sub_header).min(payload.len()) } else { 0 }; + let mut bytes = &payload[s0..]; + if skip > 0 { + let s = skip.min(bytes.len()); + skip -= s; + bytes = &bytes[s..]; + } + let take = bytes.len().min(need - out.len()); + out.extend_from_slice(&bytes[..take]); + } + } + } + pos += plen as u64; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The window refill before the zero-fill cap: always `max(n, 1 MiB)` bytes, counted when the + /// buffer has to grow. + struct OldWin<'a> { + src: &'a [u8], + buf: Vec, + start: u64, + zero_filled: usize, + } + + impl OldWin<'_> { + fn get(&mut self, pos: u64, n: usize) -> Vec { + let have = pos >= self.start && pos + n as u64 <= self.start + self.buf.len() as u64; + if !have { + let want = n.max(1 << 20); + self.zero_filled += want.saturating_sub(self.buf.len()); + self.buf.resize(want, 0); + let got = read_some(self.src, pos, &mut self.buf).unwrap(); + self.buf.truncate(got); + self.start = pos; + } + let i = (pos - self.start) as usize; + let end = (i + n).min(self.buf.len()); + self.buf[i.min(end)..end].to_vec() + } + } + + fn data(len: usize) -> Vec { + (0..len).map(|i| (i * 31 + i / 7) as u8).collect() + } + + #[test] + fn capped_window_returns_what_the_uncapped_window_returned() { + for len in [0usize, 1, 31, 32, 33, 1000, (1 << 20) - 1, 1 << 20, (1 << 20) + 5, (3 << 20) + 17] { + let d = data(len); + let mut new = Win::new(d.as_slice()); + let mut old = OldWin { src: &d, buf: Vec::new(), start: 0, zero_filled: 0 }; + let len64 = len as u64; + let mut positions = vec![ + 0, + 1, + 31, + 32, + 1 << 19, + 1 << 20, + (1 << 20) - 1, + (1 << 20) + 1, + 2 << 20, + len64.saturating_sub(40), + len64.saturating_sub(1), + len64, + len64 + 1, + len64 + 100, + ]; + positions.extend((0..10).map(|k| len64.saturating_sub(100) + k * 13)); + for &pos in &positions { + for n in [0usize, 1, 33, 1000, 65_541, 1 << 20, (1 << 20) + 1] { + assert_eq!(new.get(pos, n).unwrap(), old.get(pos, n).as_slice(), "len {len} pos {pos} n {n}"); + } + } + } + } + + #[test] + fn misses_near_the_end_zero_at_most_the_bytes_that_are_left() { + let len = (2 << 20) + 1000; + let d = data(len); + let mut new = Win::new(d.as_slice()); + let mut old = OldWin { src: &d, buf: Vec::new(), start: 0, zero_filled: 0 }; + // a PES packet whose length runs past the end of the file, asked for once per packet + let first = (2u64 << 20) as usize; + for k in 0..500 { + let pos = (first + k) as u64; + assert_eq!(new.get(pos, 2000).unwrap(), old.get(pos, 2000).as_slice()); + } + assert!(new.zero_filled <= 1000, "zeroed {} bytes, only 1000 are left in the file", new.zero_filled); + assert!(old.zero_filled > 499 * ((1 << 20) - 1000), "the reference zeroes about a megabyte per miss: {}", old.zero_filled); + } +} diff --git a/crates/media/src/mpegts/tests.rs b/crates/media/src/mpegts/tests.rs new file mode 100644 index 000000000..cec62ca90 --- /dev/null +++ b/crates/media/src/mpegts/tests.rs @@ -0,0 +1,650 @@ +//! Synthetic transport and program streams: a small muxer writes known access units, the +//! demuxer must give them back exactly (bytes, timestamps, flags); truncation and mutation never +//! panic. + +use super::*; + +fn ts_bytes(prefix: u8, v: i64) -> [u8; 5] { + let v = v & ((1 << 33) - 1); + [ + (prefix << 4) | (((v >> 30) & 7) as u8) << 1 | 1, + (v >> 22) as u8, + (((v >> 15) & 0x7F) as u8) << 1 | 1, + (v >> 7) as u8, + ((v & 0x7F) as u8) << 1 | 1, + ] +} + +pub(super) fn pes(stream_id: u8, pts: Option, dts: Option, payload: &[u8], bounded: bool) -> Vec { + let mut h = Vec::new(); + match (pts, dts) { + (Some(p), Some(d)) => { + h.extend(ts_bytes(3, p)); + h.extend(ts_bytes(1, d)); + } + (Some(p), None) => h.extend(ts_bytes(2, p)), + _ => {} + } + let flags = match (pts, dts) { + (Some(_), Some(_)) => 0xC0, + (Some(_), None) => 0x80, + _ => 0, + }; + let len = if bounded { 3 + h.len() + payload.len() } else { 0 }; + let mut v = vec![0, 0, 1, stream_id, (len >> 8) as u8, len as u8, 0x80, flags, h.len() as u8]; + v.extend(h); + v.extend_from_slice(payload); + v +} + +/// A tiny TS muxer: PAT, PMT and PES packets cut into 188-byte packets (192 with `m2ts`). +pub(super) struct TsMux { + pub(super) out: Vec, + cc: std::collections::HashMap, + m2ts: bool, +} + +impl TsMux { + pub(super) fn new(m2ts: bool) -> Self { + Self { out: Vec::new(), cc: Default::default(), m2ts } + } + fn packet(&mut self, pid: u16, pusi: bool, payload: &[u8]) -> usize { + let cc = self.cc.entry(pid).or_insert(0); + let mut p = vec![0x47, (pusi as u8) << 6 | (pid >> 8) as u8, pid as u8]; + let n = payload.len().min(184); + if n < 184 { + // adaptation field stuffing + let al = 183 - n; + p.push(0x30 | *cc); + p.push(al as u8); + if al > 0 { + p.push(0); + p.extend(std::iter::repeat_n(0xFF, al - 1)); + } + } else { + p.push(0x10 | *cc); + } + *cc = (*cc + 1) & 15; + p.extend_from_slice(&payload[..n]); + assert_eq!(p.len(), 188); + if self.m2ts { + self.out.extend_from_slice(&[0, 0, 0, 0]); + } + self.out.extend(p); + n + } + pub(super) fn section(&mut self, pid: u16, body: &[u8]) { + let mut s = body.to_vec(); + let crc = super::ts::crc32_mpeg(&s); + s.extend(crc.to_be_bytes()); + let mut payload = vec![0]; + payload.extend(s); + self.packet(pid, true, &payload); + } + pub(super) fn psi(&mut self, streams: &[(u8, u16)]) { + // PAT: program 1 → PMT PID 0x100 + let pat = [0x00, 0xB0, 13, 0, 1, 0xC1, 0, 0, 0, 1, 0xE1, 0x00]; + self.section(0, &pat); + let mut pmt = vec![0x02, 0xB0, 0, 0, 1, 0xC1, 0, 0, 0xE1, 0x00, 0xF0, 0]; + for &(st, pid) in streams { + pmt.extend([st, 0xE0 | (pid >> 8) as u8, pid as u8, 0xF0, 0]); + } + let len = pmt.len() - 3 + 4; + pmt[1] = 0xB0 | (len >> 8) as u8; + pmt[2] = len as u8; + self.section(0x100, &pmt); + } + pub(super) fn pes(&mut self, pid: u16, data: &[u8]) { + let mut rest = data; + let mut first = true; + while !rest.is_empty() { + let n = self.packet(pid, first, rest); + rest = &rest[n..]; + first = false; + } + } +} + +/// An MPEG-2 coded frame: sequence header (on I), GOP, picture header, picture coding +/// extension and a slice of filler; or a field pair. +fn mpeg2_frame(coding_type: u8, tref: u16, fields: bool, filler: usize) -> Vec { + let mut v = Vec::new(); + if coding_type == 1 { + v.extend([0, 0, 1, 0xB3, 0x14, 0x00, 0xF0, 0x13, 0xFF, 0xFF, 0xE0, 0x18]); + v.extend([0, 0, 1, 0xB8, 0, 8, 0, 0x40]); + } + let pics: &[u8] = if fields { &[1, 2] } else { &[3] }; + for &s in pics { + v.extend([0, 0, 1, 0x00, (tref >> 2) as u8, ((tref & 3) << 6) as u8 | coding_type << 3, 0xFF, 0xF8]); + v.extend([0, 0, 1, 0xB5, 0x8F, 0xFF, 0xF0 | s, 0x80, 0x80]); + v.extend([0, 0, 1, 0x01]); + v.extend((0..filler).map(|i| (i % 200) as u8 | 0x10)); + } + v +} + +/// MPEG-1 layer II frame, 48 kHz 192 kb/s (576 bytes). +fn mp2_frame(k: u8) -> Vec { + let mut v = vec![0xFF, 0xFD, 0xA4, 0x00]; + v.extend((0..572).map(|i| (i as u8) ^ k)); + v +} + +#[test] +fn transport_stream_units_round_trip() { + for m2ts in [false, true] { + let mut mux = TsMux::new(m2ts); + mux.psi(&[(0x02, 0x1011), (0x03, 0x1100)]); + // decode order I0 P3 B1 B2 with PTS = display time, DTS on anchors + let frames = [(1u8, 0u16, false), (2, 3, false), (3, 1, true), (3, 2, false)]; + let mut expect_video = Vec::new(); + for (k, &(t, tref, fields)) in frames.iter().enumerate() { + let f = mpeg2_frame(t, tref, fields, 300 + 97 * k); + let pts = 90_000 + tref as i64 * 3600; + let dts = (t != 3).then_some(90_000 - 3600 + k as i64 * 3600); + mux.pes(0x1011, &pes(0xE0, Some(pts), dts, &f, false)); + expect_video.push((f, pts, t)); + // two audio frames per PES, the second without its own timestamp + let a: Vec = [mp2_frame(2 * k as u8), mp2_frame(2 * k as u8 + 1)].concat(); + mux.pes(0x1100, &pes(0xC0, Some(90_000 + 2160 * 2 * k as i64), None, &a, true)); + } + let data = mux.out; + let f = open(&data).unwrap(); + assert_eq!(f.format, Format::Ts { packet_size: if m2ts { 192 } else { 188 }, sync_offset: if m2ts { 4 } else { 0 } }); + assert!(f.warnings.is_empty(), "{:?}", f.warnings); + assert_eq!(f.program, Some(1)); + let v = f.find(Kind::Video, |_| true).unwrap(); + let s = &f.streams[v]; + assert_eq!(s.codec, Codec::Mpeg2Video); + assert_eq!(s.units.len(), 4); + for (i, (bytes, pts, t)) in expect_video.iter().enumerate() { + let u = &s.units[i]; + assert_eq!(f.read_unit(&data, v, i).unwrap(), *bytes, "unit {i}"); + assert_eq!(u.pts, Some(*pts)); + assert_eq!(u.key, *t == 1); + assert_eq!(u.disposable, *t == 3); + assert_eq!(u.picture.unwrap().coding_type, *t); + } + assert_eq!(s.units[2].picture.unwrap().structure, 1, "field pair in one unit"); + assert!(s.units[0].picture.unwrap().sequence_header && s.units[0].picture.unwrap().gop); + let a = f.find(Kind::Audio, |c| *c == Codec::MpegAudio).unwrap(); + let s = &f.streams[a]; + assert_eq!(s.units.len(), 8); + for (i, u) in s.units.iter().enumerate() { + assert_eq!(f.read_unit(&data, a, i).unwrap(), mp2_frame(i as u8)); + assert_eq!(u.pts, (i % 2 == 0).then_some(90_000 + 2160 * i as i64)); + } + } +} + +#[test] +fn program_stream_units_round_trip() { + // packs of ≤ 2048 bytes: frames are split across PES packets + let frames: Vec> = (0..5).map(|k| mpeg2_frame(if k == 0 { 1 } else { 2 }, k as u16, k == 3, 2500 + 300 * k)).collect(); + let es: Vec = frames.concat(); + let mut starts = Vec::new(); + let mut acc = 0; + for f in &frames { + starts.push(acc); + acc += f.len(); + } + let mut out = Vec::new(); + let mut p = 0; + let mut a = 0u8; + while p < es.len() { + out.extend([0, 0, 1, 0xBA, 0x44, 0, 4, 0, 4, 1, 1, 0x89, 0xC3, 0xF8]); + let n = (es.len() - p).min(2000); + // a timestamp when a frame starts in this packet + let k = starts.iter().position(|&s| s >= p && s < p + n); + out.extend(pes(0xE0, k.map(|k| 3600 * k as i64), None, &es[p..p + n], true)); + p += n; + // AC-3 in private stream 1 (sub-stream 0x80, 4-byte header), 448 kb/s frames + let mut ac3 = vec![0x80, 1, 0, 1, 0x0B, 0x77, 0, 0, 30, 0x40]; + ac3.extend(std::iter::repeat_n(a, 1792 - 6)); + a += 1; + out.extend(pes(0xBD, Some(1000 + 2880 * a as i64), None, &ac3, true)); + } + out.extend([0, 0, 1, 0xB9]); + let f = open(&out).unwrap(); + assert_eq!(f.format, Format::Ps { mpeg1: false }); + let v = f.find(Kind::Video, |_| true).unwrap(); + assert_eq!(f.streams[v].units.len(), 5); + for (i, fr) in frames.iter().enumerate() { + assert_eq!(&f.read_unit(&out, v, i).unwrap(), fr); + assert_eq!(f.streams[v].units[i].pts, Some(3600 * i as i64)); + } + let a = f.find(Kind::Audio, |c| *c == Codec::Ac3).unwrap(); + assert_eq!(f.streams[a].id, StreamId::Ps { stream_id: 0xBD, sub_id: Some(0x80) }); + assert_eq!(f.streams[a].units.len(), a_count(&out)); + let u0 = f.read_unit(&out, a, 0).unwrap(); + assert_eq!(&u0[..2], &[0x0B, 0x77]); + assert_eq!(u0.len(), 1792); +} + +fn a_count(ps: &[u8]) -> usize { + ps.windows(4).filter(|w| *w == [0, 0, 1, 0xBD]).count() +} + +#[test] +fn sniffing() { + assert_eq!(sniff(&[0, 0, 1, 0xBA, 0x44]), Some(Format::Ps { mpeg1: false })); + assert_eq!(sniff(&[0, 0, 1, 0xBA, 0x21]), Some(Format::Ps { mpeg1: true })); + assert_eq!(sniff(&[0, 0, 1, 0xB3, 0x44]), None); + let mut ts = vec![0u8; 188 * 6]; + for k in 0..6 { + ts[k * 188] = 0x47; + } + assert!(matches!(sniff(&ts), Some(Format::Ts { packet_size: 188, .. }))); +} + +#[test] +fn truncation_and_mutation_never_panic() { + let mut mux = TsMux::new(false); + mux.psi(&[(0x02, 0x1011), (0x0F, 0x1100), (0x1B, 0x1200)]); + for k in 0..6u16 { + mux.pes(0x1011, &pes(0xE0, Some(3600 * k as i64), None, &mpeg2_frame(if k == 0 { 1 } else { 3 }, k, k % 2 == 1, 400), false)); + let mut adts = vec![0xFF, 0xF1, 0x50, 0x80, 0x0C, 0x9F, 0xFC]; + adts.extend([0x21u8; 93]); + mux.pes(0x1100, &pes(0xC0, Some(1920 * k as i64), None, &adts, true)); + mux.pes(0x1200, &pes(0xE1, Some(3600 * k as i64), None, &[0, 0, 0, 1, 0x09, 0xF0, 0, 0, 1, 0x65, 0x88, 1, 2, 3], false)); + } + let data = mux.out; + let full = open(&data).unwrap(); + assert_eq!(full.streams.len(), 3); + assert_eq!(full.streams[1].units.len(), 6, "ADTS frames"); + assert!(full.streams[2].units.iter().all(|u| u.key), "IDR access units"); + for cut in (0..data.len()).step_by(37) { + if let Ok(f) = open(&data[..cut].to_vec()) { + for (s, st) in f.streams.iter().enumerate() { + for i in 0..st.units.len() { + let _ = f.read_unit(&data[..cut].to_vec(), s, i); + } + } + } + } + let mut x = 0x9E37_79B9u32; + for _ in 0..400 { + let mut m = data.clone(); + for _ in 0..8 { + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + let i = x as usize % m.len(); + m[i] = (x >> 8) as u8; + } + if let Ok(f) = open(&m) { + for (s, st) in f.streams.iter().enumerate() { + for i in 0..st.units.len() { + let _ = f.read_unit(&m, s, i); + } + } + } + } +} + +/// Opens `data` and reads every unit; errors are fine, panics are not. +fn open_and_read_all(data: &[u8]) { + let data = data.to_vec(); + if let Ok(f) = open(&data) { + for (s, st) in f.streams.iter().enumerate() { + for i in 0..st.units.len() { + let _ = f.read_unit(&data, s, i); + } + } + } +} + +#[test] +fn program_stream_cut_inside_a_pack_header_never_panics() { + // MPEG-1 pack header cut after the 4-byte start code + 3 bytes (once panicked in `scan`). + open_and_read_all(&[0, 0, 1, 0xBA, 0x21, 0, 0]); + for n in 4..=12 { + let mpeg1 = [0u8, 0, 1, 0xBA, 0x21, 0, 1, 0, 1, 0x80, 0, 1]; + open_and_read_all(&mpeg1[..n]); + } +} + +#[test] +fn program_stream_truncation_at_every_length_never_panics() { + let frames: Vec> = (0..3).map(|k| mpeg2_frame(if k == 0 { 1 } else { 2 }, k as u16, false, 300)).collect(); + let es: Vec = frames.concat(); + for mpeg1 in [false, true] { + let mut out = Vec::new(); + for (k, chunk) in es.chunks(500).enumerate() { + if mpeg1 { + out.extend([0, 0, 1, 0xBA, 0x21, 0, 1, 0, 1, 0x80, 0, 1]); + } else { + out.extend([0, 0, 1, 0xBA, 0x44, 0, 4, 0, 4, 1, 1, 0x89, 0xC3, 0xF8]); + } + out.extend(pes(0xE0, Some(3600 * k as i64), None, chunk, true)); + } + out.extend([0, 0, 1, 0xB9]); + assert!(open(&out).is_ok(), "the untruncated stream opens (mpeg1: {mpeg1})"); + for cut in 0..=out.len() { + open_and_read_all(&out[..cut]); + } + } +} + +#[test] +fn timestamps_unwrap_across_the_33_bit_wrap() { + let mut u = Unwrap::default(); + let near = (1i64 << 33) - 3000; + assert_eq!(u.unwrap(near), near); + assert_eq!(u.unwrap(600), (1 << 33) + 600); + // a B picture timestamp from just before the wrap + assert_eq!(u.unwrap(near + 1000), near + 1000); + assert_eq!(u.unwrap(4200), (1 << 33) + 4200); +} + +/// An Annex B NAL unit: start code, header bytes and `n` bytes of payload (no zero runs). +fn nal(header: &[u8], body: &[u8]) -> Vec { + let mut v = vec![0, 0, 0, 1]; + v.extend(header); + v.extend(body); + v +} + +/// A slice of `n` filler bytes after the first byte of the slice header (0x88: first_mb 0). +fn slice(header: &[u8], first: u8, n: usize) -> Vec { + let mut body = vec![first]; + body.extend((0..n).map(|i| 0x40 | (i % 31) as u8)); + nal(header, &body) +} + +const AUD: [u8; 2] = [0x09, 0xF0]; + +fn h264_idr(n: usize) -> Vec { + [nal(&AUD, &[]), nal(&[0x67], &[0x42, 0x80, 0x1E]), nal(&[0x68], &[0xCE, 0x38, 0x80]), slice(&[0x65], 0x88, n)].concat() +} + +fn h264_p(n: usize) -> Vec { + [nal(&AUD, &[]), slice(&[0x41], 0x88, n)].concat() +} + +/// A non-reference B picture (nal_ref_idc 0). +fn h264_b(n: usize) -> Vec { + [nal(&AUD, &[]), slice(&[0x01], 0x88, n)].concat() +} + +/// A non-IDR I picture introduced by SEI with a buffering period message and a recovery point. +fn h264_recovery_i(n: usize) -> Vec { + let sei = nal(&[0x06], &[0x00, 0x02, 0xAA, 0xBB, 0x06, 0x01, 0x84, 0x80]); + [nal(&AUD, &[]), sei, slice(&[0x41], 0x88, n)].concat() +} + +fn h264_units(aus: &[(Option, Vec)]) -> (Vec, File) { + let mut mux = TsMux::new(false); + mux.psi(&[(0x1B, 0x1011)]); + for (pts, au) in aus { + mux.pes(0x1011, &pes(0xE0, *pts, None, au, false)); + } + let data = mux.out; + let f = open(&data).unwrap(); + (data, f) +} + +#[test] +fn h264_recovery_point_is_a_random_access_point() { + let aus = [(Some(90_000), h264_idr(300)), (Some(93_600), h264_p(300)), (Some(97_200), h264_recovery_i(300)), (Some(100_800), h264_p(300))]; + let (data, f) = h264_units(&aus); + let s = &f.streams[0]; + assert_eq!(s.codec, Codec::H264); + assert_eq!(s.units.iter().map(|u| u.key).collect::>(), [true, false, true, false]); + for (i, (pts, au)) in aus.iter().enumerate() { + assert_eq!(f.read_unit(&data, 0, i).unwrap(), *au, "unit {i}"); + assert_eq!(s.units[i].pts, *pts); + } + // a stream cut so that it starts at a recovery point (no IDR at all) has a key unit + let (_, f) = h264_units(&[(Some(1000), h264_recovery_i(400)), (Some(4600), h264_p(400)), (Some(8200), h264_p(400))]); + assert_eq!(f.streams[0].units.iter().map(|u| u.key).collect::>(), [true, false, false]); + // SEI without a recovery point message is not one + let sei = nal(&[0x06], &[0x00, 0x02, 0xAA, 0xBB, 0x80]); + let au = [nal(&AUD, &[]), sei, slice(&[0x41], 0x88, 600)].concat(); + let (_, f) = h264_units(&[(Some(1000), au)]); + assert!(!f.streams[0].units[0].key); +} + +#[test] +fn h264_intra_refresh_recovery_point_is_not_a_random_access_point() { + // recovery_frame_cnt 1 (ue(v) "010"): decoding from here shows broken pictures + let sei = nal(&[0x06], &[0x00, 0x02, 0xAA, 0xBB, 0x06, 0x01, 0x44, 0x80]); + let au = [nal(&AUD, &[]), sei, slice(&[0x41], 0x88, 300)].concat(); + let (_, f) = h264_units(&[(Some(1000), au), (Some(4600), h264_p(300))]); + assert_eq!(f.streams[0].units.iter().map(|u| u.key).collect::>(), [false, false]); +} + +#[test] +fn h264_several_access_units_in_one_pes_are_separate_units() { + // IDR, then in one PES a non-reference B picture and a reference P picture, only the first + // access unit of the PES has a timestamp + let (idr, b, p) = (h264_idr(400), h264_b(300), h264_p(250)); + let (data, f) = h264_units(&[(Some(90_000), idr.clone()), (Some(93_600), [b.clone(), p.clone()].concat()), (Some(100_800), h264_p(100))]); + let s = &f.streams[0]; + assert_eq!(s.units.len(), 4); + assert_eq!(s.units.iter().map(|u| u.pts).collect::>(), [Some(90_000), Some(93_600), None, Some(100_800)]); + assert_eq!(s.units.iter().map(|u| u.key).collect::>(), [true, false, false, false]); + assert_eq!(s.units.iter().map(|u| u.disposable).collect::>(), [false, true, false, false]); + assert_eq!(f.read_unit(&data, 0, 0).unwrap(), idr); + assert_eq!(f.read_unit(&data, 0, 1).unwrap(), b); + assert_eq!(f.read_unit(&data, 0, 2).unwrap(), p); +} + +#[test] +fn h264_units_without_delimiters_split_at_the_first_slice() { + // no AUD or SEI: a slice with first_mb_in_slice 0 starts the next access unit, a further + // slice of the same picture (first_mb != 0) does not + let two_slices = [slice(&[0x41], 0x88, 300), slice(&[0x41], 0x42, 300)].concat(); + let idr = slice(&[0x65], 0x88, 200); + let (data, f) = h264_units(&[(Some(0), [idr.clone(), two_slices.clone(), two_slices.clone()].concat())]); + let s = &f.streams[0]; + assert_eq!(s.units.len(), 3); + assert_eq!(s.units.iter().map(|u| u.key).collect::>(), [true, false, false]); + assert_eq!(f.read_unit(&data, 0, 1).unwrap(), two_slices); + assert_eq!(f.read_unit(&data, 0, 0).unwrap(), idr); +} + +#[test] +fn h264_one_access_unit_per_pes_is_unchanged() { + // parameter sets inside the access unit (after the AUD, before the slices) do not split it + let aus: Vec<(Option, Vec)> = + (0..6).map(|k| (Some(90_000 + 3600 * k), if k % 3 == 0 { h264_idr(500) } else { h264_p(150 + 70 * k as usize) })).collect(); + let (data, f) = h264_units(&aus); + let s = &f.streams[0]; + assert_eq!(s.units.len(), 6); + for (i, (pts, au)) in aus.iter().enumerate() { + assert_eq!(f.read_unit(&data, 0, i).unwrap(), *au); + assert_eq!(s.units[i].pts, *pts); + assert_eq!(s.units[i].key, i % 3 == 0); + assert!(!s.units[i].disposable); + } +} + +#[test] +fn hevc_several_access_units_in_one_pes_are_separate_units() { + let aud = nal(&[0x46, 0x01], &[0x50]); + let idr = [aud.clone(), slice(&[0x26, 0x01], 0x80, 300)].concat(); + // TRAIL_N picture: sub-layer non-reference + let trail_n = [aud.clone(), slice(&[0x00, 0x01], 0x80, 300)].concat(); + let trail_r = [aud, slice(&[0x02, 0x01], 0x80, 200)].concat(); + let mut mux = TsMux::new(false); + mux.psi(&[(0x24, 0x1011)]); + mux.pes(0x1011, &pes(0xE0, Some(90_000), None, &idr, false)); + mux.pes(0x1011, &pes(0xE0, Some(93_600), None, &[trail_n.clone(), trail_r.clone()].concat(), false)); + let f = open(&mux.out).unwrap(); + let s = &f.streams[0]; + assert_eq!(s.codec, Codec::Hevc); + assert_eq!(s.units.iter().map(|u| u.pts).collect::>(), [Some(90_000), Some(93_600), None]); + assert_eq!(s.units.iter().map(|u| u.key).collect::>(), [true, false, false]); + assert_eq!(s.units.iter().map(|u| u.disposable).collect::>(), [false, true, false]); + assert_eq!(f.read_unit(&mux.out, 0, 1).unwrap(), trail_n); + assert_eq!(f.read_unit(&mux.out, 0, 2).unwrap(), trail_r); +} + +fn adts_frame(k: u8) -> Vec { + let mut v = vec![0xFF, 0xF1, 0x50, 0x80, 0x0C, 0x9F, 0xFC]; + v.extend([k; 93]); + v +} + +#[test] +fn non_pes_payload_is_dropped_until_the_next_unit_start() { + // a unit start whose payload is not a PES header, then continuation packets that happen to + // hold ADTS frames: nothing of it may become a unit + let mut mux = TsMux::new(false); + mux.psi(&[(0x0F, 0x1100)]); + mux.packet(0x1100, true, &[0x55; 184]); + for k in 0..10 { + mux.packet(0x1100, false, &adts_frame(k)); + } + let f = open(&mux.out).unwrap(); + assert_eq!(f.streams.first().map(|s| s.units.len()), Some(0)); + // a real PES packet afterwards is indexed again + let mut mux2 = TsMux::new(false); + mux2.psi(&[(0x0F, 0x1100)]); + mux2.packet(0x1100, true, &[0x55; 184]); + mux2.packet(0x1100, false, &adts_frame(1)); + mux2.pes(0x1100, &pes(0xC0, Some(9000), None, &adts_frame(2), true)); + let f = open(&mux2.out).unwrap(); + let u = &f.streams[0].units; + assert_eq!(u.len(), 1); + assert_eq!(u[0].pts, Some(9000)); +} + +#[test] +fn adaptation_field_only_unit_start_keeps_the_pes_in_progress() { + let mut mux = TsMux::new(false); + mux.psi(&[(0x03, 0x1100)]); + let data = pes(0xC0, Some(9000), None, &[mp2_frame(0), mp2_frame(1)].concat(), true); + mux.packet(0x1100, true, &data[..184]); + // PUSI set, adaptation field only (no payload) + let cc = mux.cc.get(&0x1100).copied().unwrap_or(0); + let mut af = vec![0x47, 0x40 | 0x11, 0x00, 0x20 | cc, 183, 0x00]; + af.resize(188, 0xFF); + mux.out.extend(af); + for chunk in data[184..].chunks(184) { + mux.packet(0x1100, false, chunk); + } + let f = open(&mux.out).unwrap(); + let s = &f.streams[0]; + assert_eq!(s.units.len(), 2); + assert_eq!(s.units[0].pts, Some(9000)); + assert_eq!(f.read_unit(&mux.out, 0, 1).unwrap(), mp2_frame(1)); +} + +#[test] +fn oversized_unit_is_refused_without_reserving_it() { + let f = File { + format: Format::Ts { packet_size: 188, sync_offset: 0 }, + streams: vec![Stream { + id: StreamId::Pid(0x100), + stream_type: 0x1B, + codec: Codec::H264, + language: None, + units: vec![Unit { pts: None, dts: None, key: true, disposable: false, size: u32::MAX, pos: 0, offset: 0, picture: None }], + lpcm: None, + pes_packets: 0, + skipped_bytes: 0, + }], + program: None, + pcr_range: None, + warnings: Vec::new(), + }; + assert!(matches!(f.read_unit(&Vec::::new(), 0, 0), Err(Error::Invalid(_)))); +} + +/// A `ByteSource` that counts the bytes requested through `read_at`. +struct Counting<'a> { + data: &'a [u8], + read: std::cell::Cell, +} + +impl ByteSource for Counting<'_> { + fn len(&self) -> u64 { + self.data.len() as u64 + } + fn read_at(&self, offset: u64, buf: &mut [u8]) -> std::io::Result<()> { + self.read.set(self.read.get() + buf.len() as u64); + self.data.read_at(offset, buf) + } +} + +#[test] +fn bad_adaptation_fields_do_not_amplify_reads() { + // aligned packets whose adaptation_field_length (255) overruns the packet: the packet grid is + // intact, so this must cost one pass over the file, not one block read per packet + let mut data = Vec::new(); + for _ in 0..3000 { + data.extend([0x47, 0x11, 0x00, 0x30]); + data.extend(std::iter::repeat_n(0xFF, 184)); + } + let src = Counting { data: &data, read: Default::default() }; + let f = open(&src).unwrap(); + assert!(f.warnings.iter().any(|w| w.contains("adaptation field")), "{:?}", f.warnings); + assert!(!f.warnings.iter().any(|w| w.contains("lost packet sync")), "{:?}", f.warnings); + let read = src.read.get(); + assert!(read <= 4 * data.len() as u64, "read {read} bytes of a {}-byte file", data.len()); +} + +#[test] +fn garbage_without_sync_does_not_amplify_reads() { + // a sync byte every 189 bytes: five-in-a-row never lines up, every 0x47 is a false candidate + let mut data = Vec::new(); + for _ in 0..3000 { + data.extend([0x47]); + data.extend(std::iter::repeat_n(0x11, 188)); + } + let head: Vec = (0..188 * 8).map(|i| if i % 188 == 0 { 0x47 } else { 0 }).collect(); + let mut all = head; + all.extend(data); + let src = Counting { data: &all, read: Default::default() }; + let f = open(&src).unwrap(); + assert!(f.warnings.iter().any(|w| w.contains("lost packet sync")), "{:?}", f.warnings); + let read = src.read.get(); + assert!(read <= 4 * all.len() as u64, "read {read} bytes of a {}-byte file", all.len()); +} + +#[test] +fn transport_stream_resyncs_after_garbage() { + const BLOCK_PACKETS: usize = 4096; + for m2ts in [false, true] { + let ps = if m2ts { 192 } else { 188 }; + // garbage in front of the stream (shorter than 8 packets) and, at a few places around the + // first block boundary, between two frames; lengths that are not a multiple of the packet + for (delta, glen) in [(0usize, 100usize), (1, 37), (2, 188), (3, 301), (5, 5), (8, 1000)] { + let edge = (BLOCK_PACKETS - delta) * ps; + let mut mux = TsMux::new(m2ts); + mux.out.extend([0x11, 0x47, 0x22]); // garbage prefix with a lone 0x47 + mux.psi(&[(0x02, 0x1011)]); + let mut expect = Vec::new(); + let mut placed = false; + let mut k = 0usize; + while k < 3000 || !placed { + if !placed && mux.out.len() >= edge { + // garbage between two PES packets, with lone sync bytes inside + let mut g: Vec = (0..glen).map(|i| 0x10 + (i % 7) as u8).collect(); + if let Some(b) = g.get_mut(3) { + *b = 0x47; + } + mux.out.extend(g); + placed = true; + } + let f = mpeg2_frame(if k == 0 { 1 } else { 2 }, (k % 1000) as u16, false, 40 + k % 300); + mux.pes(0x1011, &pes(0xE0, Some(90_000 + 3600 * k as i64), None, &f, false)); + expect.push(f); + k += 1; + } + assert!(mux.out.len() > (BLOCK_PACKETS + 4) * ps, "stream must cross a block boundary"); + let data = mux.out; + let f = open(&data).unwrap(); + assert!(f.warnings.iter().any(|w| w.contains("lost packet sync")), "{:?}", f.warnings); + let v = f.find(Kind::Video, |_| true).unwrap(); + let s = &f.streams[v]; + assert_eq!(s.units.len(), expect.len(), "m2ts {m2ts} delta {delta} glen {glen}"); + for (i, fr) in expect.iter().enumerate() { + assert_eq!(&f.read_unit(&data, v, i).unwrap(), fr, "unit {i} (m2ts {m2ts} delta {delta} glen {glen})"); + assert_eq!(s.units[i].pts, Some(90_000 + 3600 * i as i64)); + } + } + } +} diff --git a/crates/media/src/mpegts/ts.rs b/crates/media/src/mpegts/ts.rs new file mode 100644 index 000000000..52405e14d --- /dev/null +++ b/crates/media/src/mpegts/ts.rs @@ -0,0 +1,559 @@ +//! Transport streams (H.222.0 §2.4.3): packets, PSI (PAT/PMT), PES reassembly. + +use std::collections::HashMap; + +use super::es::Splitter; +use super::pes::{self, PesHeader}; +use super::{ByteSource, Codec, Error, File, Format, Result, Stream, StreamId, Unit, Unwrap, read_some}; + +/// Byte-wise table for the MSB-first CRC-32 (polynomial 0x04C11DB7): entry `i` is the CRC register +/// after shifting the byte `i` through eight bit steps. +const CRC_TABLE: [u32; 256] = { + let mut t = [0u32; 256]; + let mut i = 0; + while i < 256 { + let mut crc = (i as u32) << 24; + let mut bit = 0; + while bit < 8 { + crc = if crc & 0x8000_0000 != 0 { (crc << 1) ^ 0x04C1_1DB7 } else { crc << 1 }; + bit += 1; + } + t[i] = crc; + i += 1; + } + t +}; + +/// CRC-32 of PSI sections (MSB-first, polynomial 0x04C11DB7, initial value all ones; Annex A). +pub(crate) fn crc32_mpeg(data: &[u8]) -> u32 { + let mut crc = 0xFFFF_FFFFu32; + for &b in data { + crc = (crc << 8) ^ CRC_TABLE[((crc >> 24) as u8 ^ b) as usize]; + } + crc +} + +/// A packet's fields. +struct Packet<'a> { + pusi: bool, + pid: u16, + cc: u8, + has_payload: bool, + discontinuity: bool, + pcr: Option, + payload: &'a [u8], +} + +/// Why a packet could not be parsed. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Bad { + /// No sync byte (or too short): the packet grid is lost. + Sync, + /// Sync byte present but the adaptation field overruns the packet: the grid is intact, the + /// packet carries nothing usable. + Adaptation, +} + +fn packet(p: &[u8]) -> std::result::Result, Bad> { + if p.len() < 188 || p[0] != 0x47 { + return Err(Bad::Sync); + } + let pusi = p[1] & 0x40 != 0; + let pid = (((p[1] & 0x1F) as u16) << 8) | p[2] as u16; + let afc = (p[3] >> 4) & 3; + let cc = p[3] & 15; + let mut start = 4; + let mut pcr = None; + let mut discontinuity = false; + if afc & 2 != 0 { + let al = p[4] as usize; + start = 5 + al; + if start > 188 { + return Err(Bad::Adaptation); + } + if al > 0 { + let flags = p[5]; + discontinuity = flags & 0x80 != 0; + if flags & 0x10 != 0 && al >= 7 { + let base = ((p[6] as i64) << 25) | ((p[7] as i64) << 17) | ((p[8] as i64) << 9) | ((p[9] as i64) << 1) | (p[10] as i64 >> 7); + let ext = (((p[10] & 1) as i64) << 8) | p[11] as i64; + pcr = Some(base * 300 + ext); + } + } + } + let has_payload = afc & 1 != 0; + Ok(Packet { pusi, pid, cc, has_payload, discontinuity, pcr, payload: if has_payload { &p[start..188] } else { &[] } }) +} + +/// Reassembles a PES header that may span packets; yields each packet's elementary-stream bytes. +#[derive(Default)] +pub(crate) struct PesAsm { + hdr: Vec, + waiting: bool, + /// The packet that started the current unit was not a PES packet: its continuation packets + /// are dropped until the next unit start. + dropping: bool, +} + +impl PesAsm { + /// The PES header completed by this packet (if any) and the packet's ES payload bytes. + pub fn packet<'a>(&mut self, pusi: bool, payload: &'a [u8]) -> (Option, &'a [u8]) { + if pusi { + if payload.is_empty() { + // an adaptation-field-only packet carries no PES data: it must not cut short + // the PES packet in progress + return (None, &[]); + } + self.hdr.clear(); + self.waiting = true; + self.dropping = false; + } else if self.dropping { + return (None, &[]); + } else if !self.waiting { + return (None, payload); + } + let before = self.hdr.len(); + self.hdr.extend_from_slice(payload); + // `pes::parse` only succeeds once the whole header (`header_len` bytes) is buffered + if let Some(h) = pes::parse(&self.hdr) { + self.waiting = false; + let skip = h.header_len.saturating_sub(before).min(payload.len()); + return (Some(h), &payload[skip..]); + } + if self.hdr.len() > 1024 || (self.hdr.len() >= 6 && !(self.hdr[0] == 0 && self.hdr[1] == 0 && self.hdr[2] == 1)) { + // not a PES packet after all: drop until the next unit start + self.waiting = false; + self.dropping = true; + self.hdr.clear(); + return (None, &[]); + } + (None, &[]) + } +} + +/// PSI section reassembly for one PID. +#[derive(Default)] +struct Psi { + buf: Vec, +} + +impl Psi { + /// Feed a packet payload; returns completed sections. + fn feed(&mut self, pusi: bool, payload: &[u8]) -> Vec> { + let mut out = Vec::new(); + if pusi { + let Some(&ptr) = payload.first() else { return out }; + let ptr = ptr as usize; + if !self.buf.is_empty() { + self.buf.extend_from_slice(payload.get(1..1 + ptr).unwrap_or(&[])); + self.drain(&mut out); + } + self.buf = payload.get(1 + ptr..).unwrap_or(&[]).to_vec(); + } else if !self.buf.is_empty() { + self.buf.extend_from_slice(payload); + } + self.drain(&mut out); + out + } + + fn drain(&mut self, out: &mut Vec>) { + while self.buf.len() >= 3 { + if self.buf[0] == 0xFF { + self.buf.clear(); + return; + } + let total = 3 + ((((self.buf[1] & 0x0F) as usize) << 8) | self.buf[2] as usize); + if self.buf.len() < total { + if total > 4096 { + self.buf.clear(); + } + return; + } + out.push(self.buf[..total].to_vec()); + self.buf.drain(..total); + } + } +} + +/// Elementary stream of the selected program. +struct EsState { + stream: Stream, + split: Splitter, + asm: PesAsm, + cc: Option, +} + +/// Map a PMT stream_type (and the ES descriptors) to a codec. `hdmv`: a Blu-ray / AVCHD stream +/// (registration descriptor "HDMV", or 192-byte packets). +pub(crate) fn codec_of(stream_type: u8, desc: &[u8], hdmv: bool) -> Codec { + let mut tags = Vec::new(); + let mut reg = None; + let mut i = 0; + while i + 2 <= desc.len() { + let (tag, len) = (desc[i], desc[i + 1] as usize); + let body = desc.get(i + 2..i + 2 + len).unwrap_or(&[]); + tags.push(tag); + if tag == 5 && body.len() >= 4 { + reg = Some([body[0], body[1], body[2], body[3]]); + } + i += 2 + len; + } + match stream_type { + 0x01 => Codec::Mpeg1Video, + 0x02 => Codec::Mpeg2Video, + 0x03 | 0x04 => Codec::MpegAudio, + 0x0F => Codec::AacAdts, + 0x11 => Codec::AacLatm, + 0x1B => Codec::H264, + 0x24 => Codec::Hevc, + 0x80 if hdmv => Codec::LpcmBluray, + 0x80 => Codec::Mpeg2Video, + 0x81 => Codec::Ac3, + 0x82 | 0x85 | 0x86 | 0xA2 if hdmv => Codec::Dts, + 0x83 if hdmv => Codec::TrueHd, + 0x84 | 0xA1 if hdmv => Codec::Eac3, + 0x87 => Codec::Eac3, + 0x90 if hdmv => Codec::Subtitle("PGS subtitles"), + 0x92 if hdmv => Codec::Subtitle("Text subtitles"), + 0x06 => { + if tags.contains(&0x6A) || reg == Some(*b"AC-3") { + Codec::Ac3 + } else if tags.contains(&0x7A) || reg == Some(*b"EAC3") { + Codec::Eac3 + } else if tags.contains(&0x7B) || matches!(reg, Some([b'D', b'T', b'S', _])) { + Codec::Dts + } else if tags.contains(&0x59) { + Codec::Subtitle("DVB subtitles") + } else if tags.contains(&0x56) { + Codec::Subtitle("Teletext") + } else { + Codec::Unknown(6) + } + } + t => Codec::Unknown(t), + } +} + +fn language(desc: &[u8]) -> Option { + let mut i = 0; + while i + 2 <= desc.len() { + let (tag, len) = (desc[i], desc[i + 1] as usize); + if tag == 0x0A && len >= 3 { + let b = desc.get(i + 2..i + 5)?; + return Some(String::from_utf8_lossy(b).trim().to_string()).filter(|s| !s.is_empty()); + } + i += 2 + len; + } + None +} + +/// First file offset at which `packet_size`-byte packets line up (sync byte at `sync_offset`). +fn first_sync(head: &[u8], packet_size: usize, sync_offset: usize) -> Option { + (0..packet_size.min(head.len())).find(|&s| (0..5).all(|k| head.get(s + sync_offset + k * packet_size) == Some(&0x47))) +} + +/// Find where packets line up again, searching `block` from byte `from`: the first packet start +/// `i` whose sync byte and the next four packets' sync bytes are all 0x47. Returns `(i, false)` +/// when found. Returns `(i, true)` when the block ran out first: `i` is where to refill from (a +/// candidate whose five-packet window does not fit yet, or the last place a whole packet fits). +/// At the end of the file (`at_end`) a candidate is accepted with as many sync bytes as exist. +fn resync(block: &[u8], from: usize, packet_size: usize, sync_offset: usize, at_end: bool) -> (usize, bool) { + let mut i = from; + while i + packet_size <= block.len() { + if block.get(i + sync_offset) == Some(&0x47) { + let mut aligned = true; + let mut complete = true; + for k in 1..5 { + match block.get(i + sync_offset + k * packet_size) { + Some(&0x47) => {} + Some(_) => { + aligned = false; + break; + } + None => { + complete = false; + break; + } + } + } + if aligned { + return (i, !complete && !at_end); + } + } + i += 1; + } + (i, true) +} + +/// PCR values unwrapped across the 33-bit base wrap, with the range seen. +#[derive(Default)] +struct PcrTrack { + unwrap: i64, + last: Option, + range: Option<(i64, i64)>, +} + +impl PcrTrack { + fn update(&mut self, pcr: i64) { + // 33-bit base × 300 wraps at 2^33·300 + let wrap = (1i64 << 33) * 300; + let mut v = pcr + self.unwrap; + if let Some(l) = self.last + && v < l - wrap / 2 + { + self.unwrap += wrap; + v += wrap; + } + self.last = Some(v); + self.range = Some(self.range.map_or((v, v), |(a, b)| (a.min(v), b.max(v)))); + } +} + +/// Register the elementary streams a (CRC-checked) PMT section announces and that are not yet known. +fn add_pmt_streams(sec: &[u8], packet_size: usize, es: &mut HashMap, order: &mut Vec) { + let pinfo = (((sec[10] & 0x0F) as usize) << 8) | sec[11] as usize; + let prog_desc = sec.get(12..12 + pinfo).unwrap_or(&[]); + let hdmv = packet_size == 192 || prog_desc.windows(4).any(|w| w == b"HDMV"); + let mut i = 12 + pinfo; + let end = sec.len() - 4; + while i + 5 <= end { + let st = sec[i]; + let pid = (((sec[i + 1] & 0x1F) as u16) << 8) | sec[i + 2] as u16; + let ilen = (((sec[i + 3] & 0x0F) as usize) << 8) | sec[i + 4] as usize; + let desc = sec.get(i + 5..(i + 5 + ilen).min(end)).unwrap_or(&[]); + i += 5 + ilen; + if es.contains_key(&pid) { + continue; + } + let mut codec = codec_of(st, desc, hdmv); + if codec == Codec::Unknown(6) && prog_desc.windows(4).any(|w| w == b"AC-3") { + codec = Codec::Ac3; + } + let stream = Stream { + id: StreamId::Pid(pid), + stream_type: st, + codec: codec.clone(), + language: language(desc), + units: Vec::new(), + lpcm: None, + pes_packets: 0, + skipped_bytes: 0, + }; + es.insert(pid, EsState { stream, split: Splitter::new(&codec), asm: PesAsm::default(), cc: None }); + order.push(pid); + } +} + +pub(crate) fn scan(src: &S, packet_size: usize, sync_offset: usize) -> Result { + let len = src.len(); + let mut head = vec![0u8; packet_size * 8]; + let n = read_some(src, 0, &mut head)?; + let mut off = first_sync(&head[..n], packet_size, sync_offset).ok_or(Error::NotMpeg)? as u64; + let mut warnings: Vec = Vec::new(); + let mut pat = Psi::default(); + let mut pmt_pid: Option = None; + let mut program = None; + let mut pmt = Psi::default(); + let mut es: HashMap = HashMap::new(); + let mut order: Vec = Vec::new(); + let mut pmt_version: Option = None; + let mut unwrap = Unwrap::default(); + let mut pcr_track = PcrTrack::default(); + let (mut lost_sync, mut cc_errors, mut crc_errors) = (0u64, 0u64, 0u64); + const BLOCK: usize = 4096; + let mut block = vec![0u8; BLOCK * packet_size]; + let mut bad_adaptation = 0u64; + while off + packet_size as u64 <= len { + let got = read_some(src, off, &mut block)?; + if got < packet_size { + break; + } + let at_end = off + got as u64 >= len; + // `p`: byte index of the next packet in `block`. Packets are parsed in place; a lost + // sync is repaired inside the loaded block, and the block is refilled from the file + // (at `off + p`) only when it is used up, so the bytes read stay close to the file size. + let mut p = 0usize; + while p + packet_size <= got { + let pos = off + p as u64; + let raw = &block[p..p + packet_size]; + let pk = match packet(&raw[sync_offset..]) { + Ok(pk) => pk, + Err(Bad::Adaptation) => { + // the packet grid is intact: this packet just has no usable payload + bad_adaptation += 1; + p += packet_size; + continue; + } + Err(Bad::Sync) => { + lost_sync += 1; + let (next, exhausted) = resync(&block[..got], p + 1, packet_size, sync_offset, at_end); + p = next; + if exhausted { + break; + } + continue; + } + }; + p += packet_size; + if let Some(pcr) = pk.pcr { + pcr_track.update(pcr); + } + if pk.pid == 0 { + for sec in pat.feed(pk.pusi, pk.payload) { + if sec.len() < 12 || sec[0] != 0 { + continue; + } + if crc32_mpeg(&sec) != 0 { + crc_errors += 1; + continue; + } + let body = &sec[8..sec.len() - 4]; + for e in body.as_chunks::<4>().0 { + let num = u16::from_be_bytes([e[0], e[1]]); + let pid = (((e[2] & 0x1F) as u16) << 8) | e[3] as u16; + if num != 0 && pmt_pid.is_none() { + pmt_pid = Some(pid); + program = Some(num); + } + } + } + continue; + } + if Some(pk.pid) == pmt_pid { + for sec in pmt.feed(pk.pusi, pk.payload) { + if sec.len() < 16 || sec[0] != 2 { + continue; + } + if crc32_mpeg(&sec) != 0 { + crc_errors += 1; + continue; + } + let version = (sec[5] >> 1) & 0x1F; + if pmt_version == Some(version) { + continue; + } + pmt_version = Some(version); + add_pmt_streams(&sec, packet_size, &mut es, &mut order); + } + continue; + } + let Some(e) = es.get_mut(&pk.pid) else { continue }; + if pk.has_payload { + if let Some(prev) = e.cc + && !pk.discontinuity + && pk.cc != (prev + 1) & 15 + && pk.cc != prev + { + cc_errors += 1; + } + e.cc = Some(pk.cc); + } + let (hdr, bytes) = e.asm.packet(pk.pusi, pk.payload); + if let Some(h) = hdr { + e.stream.pes_packets += 1; + let pts = h.pts.map(|t| unwrap.unwrap(t)); + let dts = h.dts.map(|t| unwrap.unwrap(t)); + e.split.pes_start(pts, dts); + } + e.split.feed(pos, bytes); + } + // `p >= 1` here (the loop ran at least once or `resync` moved past the failed packet) + off += p.max(1) as u64; + } + if lost_sync > 0 { + warnings.push(format!("lost packet sync {lost_sync} time(s)")); + } + if bad_adaptation > 0 { + warnings.push(format!("{bad_adaptation} packet(s) with an invalid adaptation field")); + } + if cc_errors > 0 { + warnings.push(format!("{cc_errors} continuity counter error(s)")); + } + if crc_errors > 0 { + warnings.push(format!("{crc_errors} PSI section(s) with bad CRC")); + } + if pmt_pid.is_none() { + warnings.push("no program association table".into()); + } + let mut streams = Vec::new(); + for pid in order { + let Some(mut e) = es.remove(&pid) else { continue }; + e.split.finish(); + e.stream.units = std::mem::take(&mut e.split.units); + e.stream.skipped_bytes = e.split.skipped; + if e.split.truncated { + warnings.push(format!("PID {pid}: more than {} access units, the rest are not indexed", super::es::MAX_UNITS)); + } + streams.push(e.stream); + } + Ok(File { format: Format::Ts { packet_size, sync_offset }, streams, program, pcr_range: pcr_track.range, warnings }) +} + +/// Collect unit `u` of stream `st`: from the packet at `u.pos`, the PID's elementary-stream bytes. +pub(crate) fn read(src: &S, packet_size: usize, sync_offset: usize, st: &Stream, u: &Unit, out: &mut Vec) -> Result<()> { + let StreamId::Pid(pid) = st.id else { return Err(Error::Invalid("not a TS stream".into())) }; + let need = u.size as usize; + let mut skip = u.offset as usize; + let mut off = u.pos; + let mut asm = PesAsm::default(); + let mut first = true; + let mut block = vec![0u8; 512 * packet_size]; + while out.len() < need { + let got = read_some(src, off, &mut block)?; + let whole = got / packet_size; + if whole == 0 { + break; + } + for k in 0..whole { + let raw = &block[k * packet_size..(k + 1) * packet_size]; + let pk = match packet(&raw[sync_offset..]) { + Ok(pk) => pk, + Err(Bad::Sync) => return Ok(()), + Err(Bad::Adaptation) => continue, + }; + if pk.pid != pid { + continue; + } + // a unit starting in a PES packet's first packet: the header is skipped + let (_, bytes) = if first && !pk.pusi { (None, pk.payload) } else { asm.packet(pk.pusi, pk.payload) }; + first = false; + let bytes = if skip > 0 { + let s = skip.min(bytes.len()); + skip -= s; + &bytes[s..] + } else { + bytes + }; + let take = bytes.len().min(need - out.len()); + out.extend_from_slice(&bytes[..take]); + if out.len() >= need { + return Ok(()); + } + } + off += (whole * packet_size) as u64; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A unit start that never completes a PES header (MPEG-1 syntax, stuffing without end) is + /// given up after 1024 bytes: the bytes are not kept, and the next unit start is read again. + #[test] + fn a_pes_header_that_never_ends_is_not_buffered_past_the_limit() { + let mut asm = PesAsm::default(); + let mut first = vec![0, 0, 1, 0xC0, 0, 0]; + first.resize(184, 0xFF); + assert_eq!(asm.packet(true, &first), (None, &[][..])); + let stuffing = [0xFFu8; 184]; + for k in 0..200 { + assert_eq!(asm.packet(false, &stuffing), (None, &[][..]), "packet {k}"); + assert!(asm.hdr.len() <= 1024 + 184, "{} header bytes buffered", asm.hdr.len()); + } + assert!(asm.dropping && asm.hdr.is_empty()); + let (h, es) = asm.packet(true, &[0, 0, 1, 0xC0, 0, 0, 0x80, 0, 0, 0xAA, 0xBB]); + assert_eq!((h.map(|h| h.header_len), es), (Some(9), &[0xAA, 0xBB][..])); + } +} diff --git a/crates/media/src/mpegts/work_tests.rs b/crates/media/src/mpegts/work_tests.rs new file mode 100644 index 000000000..f34aecb9f --- /dev/null +++ b/crates/media/src/mpegts/work_tests.rs @@ -0,0 +1,115 @@ +//! Equality and counting tests for the redundant-work removals: the old implementation is kept +//! here as the reference, the new one must give the same results. + +use super::*; + +/// The bitwise CRC-32 that `ts::crc32_mpeg` replaced. +fn crc32_bitwise(data: &[u8]) -> u32 { + let mut crc = 0xFFFF_FFFFu32; + for &b in data { + crc ^= (b as u32) << 24; + for _ in 0..8 { + crc = if crc & 0x8000_0000 != 0 { (crc << 1) ^ 0x04C1_1DB7 } else { crc << 1 }; + } + } + crc +} + +/// Deterministic xorshift bytes. +fn noise(seed: u32, n: usize) -> Vec { + let mut x = seed | 1; + (0..n) + .map(|_| { + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + (x >> 8) as u8 + }) + .collect() +} + +#[test] +fn table_crc_equals_the_bitwise_crc() { + // the standard check value of CRC-32/MPEG-2 + assert_eq!(ts::crc32_mpeg(b"123456789"), 0x0376_E6E7); + assert_eq!(ts::crc32_mpeg(&[]), 0xFFFF_FFFF); + for b in 0..=255u8 { + assert_eq!(ts::crc32_mpeg(&[b]), crc32_bitwise(&[b])); + } + for n in [0usize, 1, 2, 3, 4, 5, 15, 16, 17, 183, 184, 1023, 1024, 4096] { + for seed in 1..6u32 { + let d = noise(seed * 7919 + n as u32, n); + assert_eq!(ts::crc32_mpeg(&d), crc32_bitwise(&d), "len {n} seed {seed}"); + } + } + // a section followed by its own CRC checks to zero + let mut sec = noise(5, 40); + let crc = ts::crc32_mpeg(&sec); + sec.extend(crc.to_be_bytes()); + assert_eq!(ts::crc32_mpeg(&sec), 0); +} + +/// MPEG-1 layer II frame, 48 kHz 192 kb/s (576 bytes), distinct per `k` and `salt`. +fn mp2_frame(k: usize, salt: u8) -> Vec { + let mut v = vec![0xFF, 0xFD, 0xA4, 0x00]; + v.extend((0..572).map(|i| (i as u8) ^ (k as u8) ^ salt)); + v +} + +fn pes_packet(stream_id: u8, payload: &[u8]) -> Vec { + let len = 3 + payload.len(); + let mut v = vec![0, 0, 1, stream_id, (len >> 8) as u8, len as u8, 0x80, 0, 0]; + v.extend_from_slice(payload); + v +} + +/// Two audio streams interleaved in maximal PES packets over more than three window refills, +/// with a pack header per PES packet; MPEG-2 or MPEG-1 pack headers. +fn big_program_stream(mpeg1: bool, frames: usize) -> (Vec, Vec>, Vec>) { + let a: Vec> = (0..frames).map(|k| mp2_frame(k, 0)).collect(); + let b: Vec> = (0..frames).map(|k| mp2_frame(k, 0x55)).collect(); + let (ea, eb) = (a.concat(), b.concat()); + let mut out = Vec::new(); + for (ca, cb) in ea.chunks(60_000).zip(eb.chunks(60_000)) { + for (id, chunk) in [(0xC0, ca), (0xC1, cb)] { + if mpeg1 { + out.extend([0, 0, 1, 0xBA, 0x21, 0, 1, 0, 1, 0x80, 0, 1]); + } else { + out.extend([0, 0, 1, 0xBA, 0x44, 0, 4, 0, 4, 1, 1, 0x89, 0xC3, 0xF8]); + } + out.extend(pes_packet(id, chunk)); + } + } + out.extend([0, 0, 1, 0xB9]); + (out, a, b) +} + +#[test] +fn program_stream_with_large_pes_packets_gives_back_every_frame() { + for mpeg1 in [false, true] { + let (data, a, b) = big_program_stream(mpeg1, 2000); // about 1.2 MB per stream + assert!(data.len() > 2 << 20, "more than the 1 MiB window twice"); + let f = open(&data).unwrap(); + assert_eq!(f.format, Format::Ps { mpeg1 }); + assert_eq!(f.streams.len(), 2); + for (s, frames) in f.streams.iter().zip([&a, &b]) { + assert_eq!(s.units.len(), frames.len()); + // a read refills a 1 MiB window: every seventh unit and the last ones keep a debug build quick + for (i, fr) in frames.iter().enumerate().filter(|(i, _)| i % 7 == 0 || i + 4 > frames.len()) { + assert_eq!(&f.read_unit(&data, f.streams.iter().position(|x| std::ptr::eq(x, s)).unwrap(), i).unwrap(), fr, "mpeg1 {mpeg1} unit {i}"); + } + } + // cut inside the last packets: never panics, and the units before the cut are intact + for cut in [data.len() - 1, data.len() - 4, data.len() - 30_000, data.len() - 61_000, 100_000, 1 << 20] { + let d = &data[..cut]; + let Ok(f) = open(d) else { continue }; + for (si, frames) in [&a, &b].into_iter().enumerate() { + let Some(s) = f.streams.get(si) else { continue }; + let n = s.units.len().saturating_sub(1); + for (i, frame) in frames.iter().enumerate().take(n).filter(|(i, _)| i % 61 == 0 || i + 8 > n) { + assert_eq!(&f.read_unit(d, si, i).unwrap(), frame, "cut {cut} stream {si} unit {i}"); + } + } + } + } +} diff --git a/crates/media/tests/oracle_mpegts.rs b/crates/media/tests/oracle_mpegts.rs new file mode 100644 index 000000000..46a51428b --- /dev/null +++ b/crates/media/tests/oracle_mpegts.rs @@ -0,0 +1,278 @@ +//! Oracle tests: ffmpeg-made transport and program streams (skipped without ffmpeg/ffprobe). The +//! demuxer's streams, access-unit counts, sizes, timestamps and key flags are compared with +//! ffprobe's packets, and every unit reads back. + +mod common; + +use craft_media::CodecId; +use craft_media::mpegts::{Codec, File, Format, Kind, open}; +use serde_json::Value; +use std::path::{Path, PathBuf}; +use std::process::Command; + +const PIC: &[&str] = &["-f", "lavfi", "-i", "testsrc2=size=320x240:rate=25"]; +/// Stereo audio with different content per channel. +const NOISE: &[&str] = + &["-f", "lavfi", "-i", "aevalsrc=exprs=0.3*sin(2*PI*441*t)+0.1*sin(2*PI*3100*t)|0.25*sin(2*PI*660*t)+0.1*sin(2*PI*(500+4000*t)*t):s=48000"]; + +const ALL: &[&str] = &[ + "ts_mpeg2_mp2.ts", + "ts_h264_aac.ts", + "m2ts_h264_lpcm.m2ts", + "m2ts_h264_ac3.mts", + "ts_hevc_latm.ts", + "ts_h264_eac3.ts", + "ps_mpeg2_mp2.vob", + "ps_mpeg2_lpcm.vob", + "ps_mpeg2_ac3.vob", + "mpeg1_system.mpg", +]; + +fn spec(name: &str) -> Vec<&'static str> { + let cat = |p: &[&[&'static str]]| p.concat(); + let x264: &[&str] = &["-c:v", "libx264", "-bf", "2", "-g", "12", "-pix_fmt", "yuv420p"]; + match name { + // broadcast-style: interlaced MPEG-2 and MPEG-1 layer II + "ts_mpeg2_mp2.ts" => cat(&[ + &["-f", "lavfi", "-i", "testsrc2=size=720x576:rate=50,tinterlace=mode=interleave_top,setfield=tff"], + NOISE, + &[ + "-t", + "1.2", + "-c:v", + "mpeg2video", + "-flags", + "+ilme+ildct", + "-bf", + "2", + "-g", + "12", + "-b:v", + "6M", + "-c:a", + "mp2", + "-b:a", + "192k", + "-f", + "mpegts", + ], + ]), + "ts_h264_aac.ts" => cat(&[PIC, NOISE, &["-t", "1.2"], x264, &["-c:a", "aac", "-b:a", "128k", "-f", "mpegts"]]), + // BDAV (192-byte packets) with Blu-ray LPCM, and with AC-3 as AVCHD has + "m2ts_h264_lpcm.m2ts" => { + cat(&[PIC, NOISE, &["-t", "1"], x264, &["-c:a", "pcm_bluray", "-ac", "2", "-f", "mpegts", "-mpegts_m2ts_mode", "1"]]) + } + "m2ts_h264_ac3.mts" => { + cat(&[PIC, NOISE, &["-t", "1"], x264, &["-c:a", "ac3", "-b:a", "192k", "-ac", "2", "-f", "mpegts", "-mpegts_m2ts_mode", "1"]]) + } + "ts_hevc_latm.ts" => cat(&[ + PIC, + NOISE, + &["-t", "1", "-c:v", "libx265", "-x265-params", "log-level=error:bframes=2:keyint=12"], + &["-c:a", "aac", "-b:a", "96k", "-f", "mpegts", "-mpegts_flags", "latm"], + ]), + "ts_h264_eac3.ts" => cat(&[PIC, NOISE, &["-t", "1"], x264, &["-c:a", "eac3", "-b:a", "192k", "-f", "mpegts"]]), + // DVD-style program streams + "ps_mpeg2_mp2.vob" => cat(&[ + &["-f", "lavfi", "-i", "testsrc2=size=720x480:rate=30000/1001"], + NOISE, + &["-t", "1.2", "-c:v", "mpeg2video", "-bf", "2", "-g", "15", "-b:v", "5M", "-c:a", "mp2", "-b:a", "224k", "-f", "vob"], + ]), + "ps_mpeg2_lpcm.vob" => cat(&[PIC, NOISE, &["-t", "1", "-c:v", "mpeg2video", "-bf", "2", "-c:a", "pcm_dvd", "-ac", "2", "-f", "vob"]]), + "ps_mpeg2_ac3.vob" => cat(&[PIC, NOISE, &["-t", "1", "-c:v", "mpeg2video", "-bf", "2", "-c:a", "ac3", "-b:a", "192k", "-f", "vob"]]), + // ISO/IEC 11172-1 system stream + "mpeg1_system.mpg" => { + cat(&[PIC, NOISE, &["-t", "1.2", "-c:v", "mpeg1video", "-bf", "2", "-b:v", "1500k", "-c:a", "mp2", "-b:a", "128k", "-f", "mpeg"]]) + } + other => panic!("unknown fixture {other}"), + } +} + +fn fixture(name: &str) -> Option { + let ffmpeg = common::ffmpeg_or_skip("mpegts oracle")?; + let out = common::fixtures_dir("mpegts").join(name); + common::fixtures::generate(&out, |tmp: &Path| { + // the temporary name keeps no known extension: every spec names its muxer with `-f` + Command::new(&ffmpeg).args(["-v", "error", "-y", "-nostdin"]).args(spec(name)).arg(tmp).status().is_ok_and(|s| s.success()) + }) +} + +/// ffprobe's packets of stream `index`: (pts, size, key). +fn packets(all: &[Value], index: i64) -> Vec<(Option, u64, bool)> { + all.iter() + .filter(|p| common::int(&p["stream_index"]) == Some(index)) + .map(|p| (common::int(&p["pts"]), common::int(&p["size"]).unwrap_or(0) as u64, p["flags"].as_str().is_some_and(|f| f.starts_with('K')))) + .collect() +} + +/// The codec ffprobe names, as this module's `Codec`. +fn expected_codec(name: &str, format: Format) -> Codec { + match name { + "mpeg1video" => Codec::Mpeg1Video, + "mpeg2video" => Codec::Mpeg2Video, + "h264" => Codec::H264, + "hevc" => Codec::Hevc, + "mp2" | "mp3" => Codec::MpegAudio, + "aac" => Codec::AacAdts, + "aac_latm" => Codec::AacLatm, + "ac3" => Codec::Ac3, + "eac3" => Codec::Eac3, + "pcm_bluray" => Codec::LpcmBluray, + "pcm_dvd" => Codec::LpcmDvd, + other => panic!("{format:?}: ffprobe codec {other} has no mapping in this test"), + } +} + +fn check(name: &str) -> Option { + let path = fixture(name)?; + let bytes = std::fs::read(&path).unwrap(); + let f = open(&bytes).unwrap(); + assert!(f.warnings.is_empty(), "{name}: {:?}", f.warnings); + let streams = common::streams(&path)?; + let all = common::packets(&path)?; + assert_eq!(f.streams.len(), streams.len(), "{name}: stream count"); + for kind in [Kind::Video, Kind::Audio] { + let what = if kind == Kind::Video { "video" } else { "audio" }; + // a program stream lists its streams in the order they appear, which need not be ffprobe's + let ours = f.streams.iter().position(|s| s.kind() == kind).unwrap_or_else(|| panic!("{name}: no {what} stream")); + let theirs = streams.iter().find(|s| s["codec_type"] == what).unwrap_or_else(|| panic!("{name}: ffprobe has no {what} stream")); + let s = &f.streams[ours]; + let codec = expected_codec(theirs["codec_name"].as_str().unwrap(), f.format); + assert_eq!(s.codec, codec, "{name}: {what} codec"); + let pk = packets(&all, common::int(&theirs["index"]).unwrap()); + assert_eq!(s.units.len(), pk.len(), "{name}: {what} unit count"); + if kind == Kind::Video { + // every picture ffprobe gives a PTS has the same one here (ffprobe fills in some a + // program stream leaves out), in the same order + for (i, (u, p)) in s.units.iter().zip(&pk).enumerate() { + if let Some(pts) = u.pts { + assert_eq!(Some(pts), p.0, "{name}: video unit {i} PTS"); + } + assert_eq!(u.key, p.2, "{name}: video unit {i} key flag"); + } + let stamped = s.units.iter().filter(|u| u.pts.is_some()).count(); + assert!(stamped > 0, "{name}: no video unit has a PTS"); + if matches!(f.format, Format::Ts { .. }) { + assert_eq!(stamped, s.units.len(), "{name}: every TS picture starts a PES packet"); + } + } else { + // LATM: ffprobe's packets are the AAC frames without the LOAS framing + if codec != Codec::AacLatm { + let (ours, theirs): (Vec, Vec) = (s.units.iter().map(|u| u64::from(u.size)).collect(), pk.iter().map(|p| p.1).collect()); + // a DVD LPCM unit is the samples alone; ffprobe's packet keeps the 3 format bytes + // before them + let kept = if codec == Codec::LpcmDvd { 3 } else { 0 }; + assert_eq!(ours.iter().map(|n| n + kept).collect::>(), theirs, "{name}: audio unit sizes"); + } + // a PTS belongs to the first frame that starts in a PES packet; ffprobe interpolates + // the others + for (i, (u, p)) in s.units.iter().zip(&pk).enumerate() { + if let Some(pts) = u.pts { + assert_eq!(Some(pts), p.0, "{name}: audio unit {i} PTS"); + } + } + assert!(s.units.first().is_some_and(|u| u.pts.is_some()), "{name}: the first audio unit has a PTS"); + assert!(s.units.iter().all(|u| u.key), "{name}: audio units are all key"); + } + // every unit reads back whole + for (i, u) in s.units.iter().enumerate() { + assert_eq!(f.read_unit(&bytes, ours, i).unwrap().len(), u.size as usize, "{name}: {what} unit {i}"); + } + } + Some(f) +} + +#[test] +fn transport_streams_match_ffprobe() { + for name in ["ts_mpeg2_mp2.ts", "ts_h264_aac.ts", "ts_hevc_latm.ts", "ts_h264_eac3.ts"] { + let Some(f) = check(name) else { return }; + assert_eq!(f.format, Format::Ts { packet_size: 188, sync_offset: 0 }, "{name}"); + assert_eq!(f.program, Some(1), "{name}"); + let (first, last) = f.pcr_range.unwrap(); + assert!(last > first, "{name}: PCR range"); + let ids: Vec = f.streams.iter().map(|s| s.codec_params().codec).collect(); + let want = match name { + "ts_mpeg2_mp2.ts" => [CodecId::Mpeg2Video, CodecId::Mp2], + "ts_h264_aac.ts" => [CodecId::H264, CodecId::Aac], + "ts_hevc_latm.ts" => [CodecId::Hevc, CodecId::Aac], + _ => [CodecId::H264, CodecId::Eac3], + }; + assert_eq!(ids, want, "{name}"); + } +} + +#[test] +fn mpeg2_pictures_carry_their_coding_type_and_field_order() { + let Some(f) = check("ts_mpeg2_mp2.ts") else { return }; + let v = &f.streams[0]; + let pics: Vec<_> = v.units.iter().map(|u| u.picture.unwrap()).collect(); + assert_eq!(pics[0].coding_type, 1); + assert!(pics[0].sequence_header && pics[0].gop); + assert!(pics.iter().any(|p| p.coding_type == 3) && pics.iter().any(|p| p.coding_type == 2)); + assert!(pics.iter().all(|p| p.top_field_first && !p.progressive_frame), "interlaced, top field first"); + // B pictures are the disposable ones, I pictures the keys + for (u, p) in v.units.iter().zip(&pics) { + assert_eq!((u.key, u.disposable), (p.coding_type == 1, p.coding_type == 3)); + } +} + +#[test] +fn bdav_streams_match_ffprobe() { + for name in ["m2ts_h264_lpcm.m2ts", "m2ts_h264_ac3.mts"] { + let Some(f) = check(name) else { return }; + assert_eq!(f.format, Format::Ts { packet_size: 192, sync_offset: 4 }, "{name}"); + assert_eq!(f.format.name(), "MPEG-2 TS (BDAV/AVCHD)"); + } +} + +#[test] +fn program_streams_match_ffprobe() { + for name in ["ps_mpeg2_mp2.vob", "ps_mpeg2_lpcm.vob", "ps_mpeg2_ac3.vob", "mpeg1_system.mpg"] { + let Some(f) = check(name) else { return }; + assert_eq!(f.format, Format::Ps { mpeg1: name == "mpeg1_system.mpg" }, "{name}"); + assert!(f.pcr_range.is_some(), "{name}: SCR range"); + if name == "ps_mpeg2_lpcm.vob" { + let a = f.streams.iter().find(|s| s.kind() == Kind::Audio).unwrap(); + let p = a.codec_params(); + assert_eq!((p.codec, p.sample_rate, p.channels), (CodecId::Pcm, Some(48_000), Some(2))); + assert!(p.pcm.is_some_and(|f| f.big_endian && f.signed && !f.float)); + } + } +} + +/// Every fixture, truncated and damaged: opening and reading never panics. +#[test] +fn damaged_ffmpeg_files_never_panic() { + for name in ALL { + let Some(path) = fixture(name) else { return }; + let bytes = std::fs::read(&path).unwrap(); + let mut x = 0x2545_F491u32; + let mut next = move || { + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + x as usize + }; + for round in 0..12 { + let mut d = bytes.clone(); + if round % 3 == 0 { + d.truncate(next() % d.len()); + } + for _ in 0..(round * 40) { + let at = next() % d.len().max(1); + if let Some(b) = d.get_mut(at) { + *b = next() as u8; + } + } + let Ok(f) = open(&d) else { continue }; + for (s, st) in f.streams.iter().enumerate() { + let _ = st.codec_params(); + // the first, the last and a spread of units + let n = st.units.len(); + for i in (0..n).step_by((n / 16).max(1)).chain(n.checked_sub(1)) { + let _ = f.read_unit(&d, s, i); + } + } + } + } +}