use super::mkv::{MkvMuxer, MkvTrack};
use super::{WriteSeek, ebml};
type MkvHeaderResult = io::Result<(crate::disc::DiscTitle, Vec<(u16, Vec<u8>)>, i64)>;
fn skip_bytes(r: &mut impl Read, n: u64) -> io::Result<()> {
io::copy(&mut r.take(n), &mut io::sink())?;
Ok(())
}
const MAX_BLOCK_SIZE: u64 = 64 * 1024 * 1024;
const MAX_CODEC_PRIVATE: u64 = 16 * 1024 * 1024;
const MAX_STRING_LEN: u64 = 64 * 1024;
const MAX_UINT_LEN: u64 = 8;
fn checked_size(size: u64, cap: u64) -> io::Result<usize> {
if size > cap {
return Err(crate::error::Error::MkvInvalid.into());
}
Ok(size as usize)
}
fn read_uint_bounded(r: &mut impl Read, size: u64) -> io::Result<u64> {
ebml::read_uint_val(r, checked_size(size, MAX_UINT_LEN)?)
}
fn read_string_bounded(r: &mut impl Read, size: u64) -> io::Result<String> {
ebml::read_string_val(r, checked_size(size, MAX_STRING_LEN)?)
}
use crate::disc::*;
use std::io::{self, Read};
struct ReadState {
reader: Box<dyn Read + Send>,
cluster_ts_ticks: i64,
ts_scale_ns: i64,
codec_privates: Vec<(u16, Vec<u8>)>,
}
enum Mode {
Write {
muxer: Option<MkvMuxer<Box<dyn WriteSeek + Send>>>,
},
Read(ReadState),
}
pub struct MkvStream {
disc_title: DiscTitle,
mode: Mode,
}
impl MkvStream {
pub fn create(writer: Box<dyn WriteSeek + Send>, title: &DiscTitle) -> io::Result<Self> {
let mut tracks = Vec::new();
let mut has_default_video = false;
let mut has_default_audio = false;
for (idx, s) in title.streams.iter().enumerate() {
let mut track = match s {
crate::disc::Stream::Video(v) => MkvTrack::video(v),
crate::disc::Stream::Audio(a) => MkvTrack::audio(a),
crate::disc::Stream::Subtitle(s) => MkvTrack::subtitle(s),
};
if track.is_default {
match track.track_type {
1 if !has_default_video => has_default_video = true,
2 if !has_default_audio => has_default_audio = true,
_ => track.is_default = false,
}
}
if let Some(cp) = title.codec_privates.get(idx).and_then(|c| c.as_ref()) {
track.codec_private = Some(cp.clone());
}
tracks.push(track);
}
let muxer = MkvMuxer::new(
writer,
&tracks,
Some(&title.playlist),
title.duration_secs,
&title.chapters,
)?;
Ok(Self {
disc_title: title.clone(),
mode: Mode::Write { muxer: Some(muxer) },
})
}
pub fn open(mut reader: impl Read + Send + 'static) -> io::Result<Self> {
let (disc_title, codec_privates, ts_scale_ns) = parse_mkv_header(&mut reader)?;
Ok(Self {
disc_title,
mode: Mode::Read(ReadState {
reader: Box::new(reader),
cluster_ts_ticks: 0,
ts_scale_ns,
codec_privates,
}),
})
}
}
impl crate::pes::Stream for MkvStream {
fn read(&mut self) -> io::Result<Option<crate::pes::PesFrame>> {
let streams_len = self.disc_title.streams.len();
let rs = match self.mode {
Mode::Read(ref mut rs) => rs,
Mode::Write { .. } => return Err(crate::error::Error::StreamWriteOnly.into()),
};
loop {
let (id, size, _) = match ebml::read_element_header(&mut rs.reader) {
Ok(h) => h,
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
Err(e) => return Err(e),
};
match id {
ebml::CLUSTER => continue,
ebml::CLUSTER_TIMESTAMP => {
let raw = read_uint_bounded(&mut rs.reader, size)?;
if raw > i64::MAX as u64 {
return Err(crate::error::Error::MkvInvalid.into());
}
rs.cluster_ts_ticks = raw as i64;
continue;
}
ebml::SIMPLE_BLOCK => {
let block =
ebml::read_binary_val(&mut rs.reader, checked_size(size, MAX_BLOCK_SIZE)?)?;
if let Some(frame) = parse_block(
&block,
rs.cluster_ts_ticks,
rs.ts_scale_ns,
streams_len,
None,
) {
return Ok(Some(frame));
}
continue;
}
ebml::BLOCK_GROUP => {
if size == u64::MAX {
return Err(crate::error::Error::MkvInvalid.into());
}
let mut remaining = size;
let mut block: Option<Vec<u8>> = None;
let mut duration_ms: Option<u64> = None;
while remaining > 0 {
let (cid, cs, hlen) = ebml::read_element_header(&mut rs.reader)?;
if cs == u64::MAX {
return Err(crate::error::Error::MkvInvalid.into());
}
remaining = remaining.saturating_sub(hlen as u64 + cs);
match cid {
ebml::BLOCK => {
block = Some(ebml::read_binary_val(
&mut rs.reader,
checked_size(cs, MAX_BLOCK_SIZE)?,
)?);
}
ebml::BLOCK_DURATION => {
duration_ms = Some(read_uint_bounded(&mut rs.reader, cs)?);
}
_ => skip_bytes(&mut rs.reader, cs)?,
}
}
if let Some(block) = block {
let dur_ns = duration_ms.map(|ms| ms.saturating_mul(1_000_000));
if let Some(frame) = parse_block(
&block,
rs.cluster_ts_ticks,
rs.ts_scale_ns,
streams_len,
dur_ns,
) {
return Ok(Some(frame));
}
}
continue;
}
_ => {
if size == u64::MAX {
return Err(crate::error::Error::MkvInvalid.into());
}
skip_bytes(&mut rs.reader, size)?;
continue;
}
}
}
}
fn write(&mut self, frame: &crate::pes::PesFrame) -> io::Result<()> {
match &mut self.mode {
Mode::Write { muxer: Some(m) } => m.write_frame(
frame.track,
frame.pts,
frame.keyframe,
&frame.data,
frame.duration_ns,
),
Mode::Write { muxer: None } => Ok(()),
Mode::Read(_) => Err(crate::error::Error::StreamReadOnly.into()),
}
}
fn finish(&mut self) -> io::Result<()> {
if let Mode::Write { ref mut muxer } = self.mode {
if let Some(m) = muxer.take() {
m.finish()?;
}
}
Ok(())
}
fn info(&self) -> &crate::disc::DiscTitle {
&self.disc_title
}
fn codec_private(&self, track: usize) -> Option<Vec<u8>> {
let track_num = (track + 1) as u16; if let Mode::Read(ref rs) = self.mode {
rs.codec_privates
.iter()
.find(|(tn, _)| *tn == track_num)
.map(|(_, data)| data.clone())
} else {
None
}
}
fn headers_ready(&self) -> bool {
true }
}
fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult {
let mut title = String::new();
let mut duration_ticks = 0.0f64;
let mut ts_scale: u64 = 1_000_000;
let mut streams: Vec<crate::disc::Stream> = Vec::new();
let mut codec_privates: Vec<(u16, Vec<u8>)> = Vec::new();
let (id, size, _) = ebml::read_element_header(r)?;
if id != ebml::EBML {
return Err(crate::error::Error::MkvInvalid.into());
}
if size > i64::MAX as u64 {
return Err(crate::error::Error::MkvInvalid.into());
}
skip_bytes(r, size)?;
let (id, _, _) = ebml::read_element_header(r)?;
if id != ebml::SEGMENT {
return Err(crate::error::Error::MkvInvalid.into());
}
let (mut got_info, mut got_tracks) = (false, false);
loop {
if got_info && got_tracks {
break;
}
let (id, size, _) = match ebml::read_element_header(r) {
Ok(h) => h,
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
Err(e) => return Err(e),
};
match id {
ebml::INFO => {
if size == u64::MAX {
return Err(crate::error::Error::MkvInvalid.into());
}
let mut remaining = size;
while remaining > 0 {
let (cid, cs, hlen) = ebml::read_element_header(r)?;
if cs == u64::MAX {
return Err(crate::error::Error::MkvInvalid.into());
}
remaining = remaining.saturating_sub(hlen as u64 + cs);
match cid {
ebml::TIMESTAMP_SCALE => ts_scale = read_uint_bounded(r, cs)?,
ebml::DURATION => duration_ticks = ebml::read_float_val(r, cs as usize)?,
ebml::TITLE => title = read_string_bounded(r, cs)?,
_ => {
skip_bytes(r, cs)?;
}
}
}
got_info = true;
}
ebml::TRACKS => {
if size == u64::MAX {
return Err(crate::error::Error::MkvInvalid.into());
}
let mut remaining = size;
while remaining > 0 {
let (cid, cs, hlen) = ebml::read_element_header(r)?;
if cs == u64::MAX {
return Err(crate::error::Error::MkvInvalid.into());
}
remaining = remaining.saturating_sub(hlen as u64 + cs);
if cid == ebml::TRACK_ENTRY {
let (stream, tnum, cp) = parse_track(r, cs)?;
if let Some(s) = stream {
streams.push(s);
}
if let Some(cp) = cp {
codec_privates.push((tnum, cp));
}
} else {
skip_bytes(r, cs)?;
}
}
got_tracks = true;
}
ebml::CLUSTER => break,
_ if size != u64::MAX => {
skip_bytes(r, size)?;
}
_ => break,
}
}
let disc_title = DiscTitle {
playlist: title,
duration_secs: duration_ticks * (ts_scale as f64) / 1_000_000_000.0,
streams,
..DiscTitle::empty()
};
let ts_scale_ns = if ts_scale == 0 || ts_scale > i64::MAX as u64 {
1_000_000
} else {
ts_scale as i64
};
Ok((disc_title, codec_privates, ts_scale_ns))
}
const MAX_TS_PID: u32 = 0x1FFF;
fn ts_pid_for_track(tnum: u16) -> io::Result<u16> {
if tnum == 0 {
return Err(crate::error::Error::MkvInvalid.into());
}
let pid: u32 = if tnum == 1 {
0x1011
} else {
0x1100u32 + (tnum as u32 - 2)
};
if pid > MAX_TS_PID {
return Err(crate::error::Error::MkvInvalid.into());
}
Ok(pid as u16)
}
fn parse_track(
r: &mut impl Read,
size: u64,
) -> io::Result<(Option<crate::disc::Stream>, u16, Option<Vec<u8>>)> {
let (mut ttype, mut tnum) = (0u64, 0u16);
let (mut codec_id, mut lang, mut name) = (String::new(), String::from("und"), String::new());
let (mut ph, mut sr, mut ch, mut forced) = (0u32, 0.0f64, 0u8, false);
let mut codec_priv: Option<Vec<u8>> = None;
let mut remaining = size;
while remaining > 0 {
let (cid, cs, hlen) = ebml::read_element_header(r)?;
if cs == u64::MAX {
return Err(crate::error::Error::MkvInvalid.into());
}
remaining = remaining.saturating_sub(hlen as u64 + cs);
match cid {
ebml::TRACK_NUMBER => {
let n = read_uint_bounded(r, cs)?;
if n > u16::MAX as u64 {
return Err(crate::error::Error::MkvInvalid.into());
}
tnum = n as u16;
}
ebml::TRACK_TYPE => ttype = read_uint_bounded(r, cs)?,
ebml::CODEC_ID => codec_id = read_string_bounded(r, cs)?,
ebml::CODEC_PRIVATE => {
codec_priv = Some(ebml::read_binary_val(
r,
checked_size(cs, MAX_CODEC_PRIVATE)?,
)?)
}
ebml::LANGUAGE => lang = read_string_bounded(r, cs)?,
ebml::TRACK_NAME => name = read_string_bounded(r, cs)?,
ebml::FLAG_FORCED => forced = read_uint_bounded(r, cs)? != 0,
ebml::VIDEO => {
let mut vrem = cs;
while vrem > 0 {
let (vid, vs, vhlen) = ebml::read_element_header(r)?;
if vs == u64::MAX {
return Err(crate::error::Error::MkvInvalid.into());
}
vrem = vrem.saturating_sub(vhlen as u64 + vs);
if vid == ebml::PIXEL_HEIGHT {
ph = read_uint_bounded(r, vs)? as u32;
} else {
skip_bytes(r, vs)?;
}
}
}
ebml::AUDIO => {
let mut arem = cs;
while arem > 0 {
let (aid, as_, ahlen) = ebml::read_element_header(r)?;
if as_ == u64::MAX {
return Err(crate::error::Error::MkvInvalid.into());
}
arem = arem.saturating_sub(ahlen as u64 + as_);
match aid {
ebml::SAMPLING_FREQUENCY => sr = ebml::read_float_val(r, as_ as usize)?,
ebml::CHANNELS => ch = read_uint_bounded(r, as_)? as u8,
_ => {
skip_bytes(r, as_)?;
}
}
}
}
_ => {
skip_bytes(r, cs)?;
}
}
}
let codec = match codec_id.as_str() {
"V_MPEGH/ISO/HEVC" => Codec::Hevc,
"V_MPEG4/ISO/AVC" => Codec::H264,
"V_MS/VFW/FOURCC" => Codec::Vc1,
"V_MPEG2" => Codec::Mpeg2,
"A_AC3" => Codec::Ac3,
"A_EAC3" => Codec::Ac3Plus,
"A_TRUEHD" => Codec::TrueHd,
"A_DTS" => Codec::Dts,
"A_PCM/INT/BIG" => Codec::Lpcm,
"S_HDMV/PGS" => Codec::Pgs,
"S_VOBSUB" => Codec::DvdSub,
_ => Codec::Unknown(0),
};
let res = Resolution::from_height(ph);
let chs = AudioChannels::from_count(ch);
let srs = if sr >= 96000.0 {
SampleRate::S96
} else {
SampleRate::S48
};
let ts_pid = ts_pid_for_track(tnum)?;
let stream = match ttype {
1 => {
let is_secondary = name.contains("Dolby Vision EL") || name.contains("DV EL");
Some(crate::disc::Stream::Video(VideoStream {
pid: ts_pid,
codec,
resolution: res,
frame_rate: FrameRate::Unknown,
hdr: HdrFormat::Sdr,
color_space: ColorSpace::Bt709,
secondary: is_secondary,
label: name,
}))
}
2 => Some(crate::disc::Stream::Audio(AudioStream {
pid: ts_pid,
codec,
channels: chs,
language: lang,
sample_rate: srs,
secondary: false,
purpose: crate::disc::LabelPurpose::Normal,
label: name,
})),
17 => Some(crate::disc::Stream::Subtitle(SubtitleStream {
pid: ts_pid,
codec,
language: lang,
forced,
qualifier: crate::disc::LabelQualifier::None,
codec_data: None,
})),
_ => None,
};
Ok((stream, tnum, codec_priv))
}
fn parse_block(
block: &[u8],
cluster_ts_ticks: i64,
ts_scale_ns: i64,
streams_len: usize,
duration_ns: Option<u64>,
) -> Option<crate::pes::PesFrame> {
if block.len() < 4 {
return None;
}
let (track, vl) = block_vint(block);
if vl + 3 > block.len() {
return None;
}
if track == 0 {
return None;
}
let rel_ts = i16::from_be_bytes([block[vl], block[vl + 1]]);
let keyframe = block[vl + 2] & 0x80 != 0;
let data = block[vl + 3..].to_vec();
let pts_ticks = cluster_ts_ticks + rel_ts as i64;
let track_idx = (track as usize) - 1;
if track_idx >= streams_len {
return None;
}
Some(crate::pes::PesFrame {
track: track_idx,
pts: pts_ticks.saturating_mul(ts_scale_ns),
keyframe,
data,
duration_ns,
})
}
fn block_vint(d: &[u8]) -> (u64, usize) {
if d.is_empty() {
return (0, 0);
}
if d[0] & 0x80 != 0 {
return ((d[0] & 0x7F) as u64, 1);
}
if d[0] & 0x40 != 0 && d.len() >= 2 {
return ((((d[0] & 0x3F) as u64) << 8) | d[1] as u64, 2);
}
if d[0] & 0x20 != 0 && d.len() >= 3 {
return (
(((d[0] & 0x1F) as u64) << 16) | ((d[1] as u64) << 8) | d[2] as u64,
3,
);
}
if d[0] & 0x10 != 0 && d.len() >= 4 {
return (
(((d[0] & 0x0F) as u64) << 24)
| ((d[1] as u64) << 16)
| ((d[2] as u64) << 8)
| d[3] as u64,
4,
);
}
(0, 1) }
#[cfg(test)]
mod tests {
use super::*;
use crate::pes::Stream as _;
use std::io::Cursor;
fn open_err(r: io::Result<MkvStream>) -> io::Error {
match r {
Ok(_) => panic!("expected MkvStream::open to fail"),
Err(e) => e,
}
}
fn is_mkv_invalid(e: &io::Error) -> bool {
e.kind() == io::ErrorKind::InvalidData
&& e.to_string()
.starts_with(&format!("E{}", crate::error::E_MKV_INVALID))
}
#[test]
fn ts_pid_for_track_maps_and_rejects_overflow() {
assert_eq!(ts_pid_for_track(1).unwrap(), 0x1011);
assert_eq!(ts_pid_for_track(2).unwrap(), 0x1100);
assert_eq!(ts_pid_for_track(3).unwrap(), 0x1101);
assert_eq!(ts_pid_for_track(0xF01).unwrap(), 0x1FFF);
assert!(is_mkv_invalid(&ts_pid_for_track(0xF02).unwrap_err()));
assert!(is_mkv_invalid(&ts_pid_for_track(u16::MAX).unwrap_err()));
assert!(is_mkv_invalid(&ts_pid_for_track(0).unwrap_err()));
}
#[test]
fn checked_size_rejects_over_cap() {
assert_eq!(checked_size(100, 256).unwrap(), 100);
assert_eq!(checked_size(256, 256).unwrap(), 256);
let e = checked_size(257, 256).unwrap_err();
assert!(is_mkv_invalid(&e));
let e = checked_size(4 * 1024 * 1024 * 1024, MAX_BLOCK_SIZE).unwrap_err();
assert!(is_mkv_invalid(&e));
}
#[test]
fn read_uint_bounded_rejects_oversized_int() {
let mut data = Cursor::new(vec![0u8; 16]);
let e = read_uint_bounded(&mut data, 9).unwrap_err();
assert!(is_mkv_invalid(&e));
}
#[test]
fn read_uint_bounded_accepts_valid_width() {
let mut data = Cursor::new(vec![0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x02]);
assert_eq!(read_uint_bounded(&mut data, 8).unwrap(), 0x0102);
}
#[test]
fn read_string_bounded_rejects_huge_string() {
let mut data = Cursor::new(vec![0u8; 16]);
let e = read_string_bounded(&mut data, MAX_STRING_LEN + 1).unwrap_err();
assert!(is_mkv_invalid(&e));
}
fn minimal_mkv_with_cluster(cluster_body: &[u8]) -> Vec<u8> {
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::TRACKS).unwrap();
ebml::write_size(&mut out, 0).unwrap();
out.extend_from_slice(cluster_body);
out
}
#[test]
fn simple_block_oversized_size_is_rejected() {
let mut cluster = Vec::new();
ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap();
ebml::write_unknown_size(&mut cluster).unwrap();
ebml::write_id(&mut cluster, ebml::SIMPLE_BLOCK).unwrap();
ebml::write_size(&mut cluster, 2 * 1024 * 1024 * 1024).unwrap();
let bytes = minimal_mkv_with_cluster(&cluster);
let mut stream = MkvStream::open(Cursor::new(bytes)).unwrap();
let e = stream.read().unwrap_err();
assert!(is_mkv_invalid(&e));
}
#[test]
fn well_formed_simple_block_round_trips() {
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, 0).unwrap();
let mut entry = Vec::new();
ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, 1).unwrap();
ebml::write_uint(&mut entry, ebml::TRACK_TYPE, 1).unwrap();
let mut track_entry = Vec::new();
ebml::write_id(&mut track_entry, ebml::TRACK_ENTRY).unwrap();
ebml::write_size(&mut track_entry, entry.len() as u64).unwrap();
track_entry.extend_from_slice(&entry);
ebml::write_id(&mut out, ebml::TRACKS).unwrap();
ebml::write_size(&mut out, track_entry.len() as u64).unwrap();
out.extend_from_slice(&track_entry);
ebml::write_id(&mut out, ebml::CLUSTER).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
let block = [0x81u8, 0x00, 0x00, 0x80, 0xAA, 0xBB, 0xCC, 0xDD];
ebml::write_id(&mut out, ebml::SIMPLE_BLOCK).unwrap();
ebml::write_size(&mut out, block.len() as u64).unwrap();
out.extend_from_slice(&block);
let mut stream = MkvStream::open(Cursor::new(out)).unwrap();
let frame = stream.read().unwrap().expect("expected a frame");
assert_eq!(frame.track, 0);
assert!(frame.keyframe);
assert_eq!(frame.data, vec![0xAA, 0xBB, 0xCC, 0xDD]);
}
#[test]
fn truncated_simple_block_body_errors_not_panics() {
let mut cluster = Vec::new();
ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap();
ebml::write_unknown_size(&mut cluster).unwrap();
ebml::write_id(&mut cluster, ebml::SIMPLE_BLOCK).unwrap();
ebml::write_size(&mut cluster, 64).unwrap();
let bytes = minimal_mkv_with_cluster(&cluster);
let mut stream = MkvStream::open(Cursor::new(bytes)).unwrap();
let e = stream.read().unwrap_err();
assert!(is_mkv_invalid(&e));
}
fn mkv_with_track_and_cluster(tnum: u64, ttype: u64, cluster_body: &[u8]) -> Vec<u8> {
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, 0).unwrap();
let mut entry = Vec::new();
ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, tnum).unwrap();
ebml::write_uint(&mut entry, ebml::TRACK_TYPE, ttype).unwrap();
let mut track_entry = Vec::new();
ebml::write_id(&mut track_entry, ebml::TRACK_ENTRY).unwrap();
ebml::write_size(&mut track_entry, entry.len() as u64).unwrap();
track_entry.extend_from_slice(&entry);
ebml::write_id(&mut out, ebml::TRACKS).unwrap();
ebml::write_size(&mut out, track_entry.len() as u64).unwrap();
out.extend_from_slice(&track_entry);
out.extend_from_slice(cluster_body);
out
}
#[test]
fn oversized_codec_private_is_rejected() {
let mut entry = Vec::new();
ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, 1).unwrap();
ebml::write_uint(&mut entry, ebml::TRACK_TYPE, 1).unwrap();
ebml::write_id(&mut entry, ebml::CODEC_PRIVATE).unwrap();
ebml::write_size(&mut entry, MAX_CODEC_PRIVATE + 1).unwrap();
let mut track_entry = Vec::new();
ebml::write_id(&mut track_entry, ebml::TRACK_ENTRY).unwrap();
ebml::write_size(&mut track_entry, entry.len() as u64).unwrap();
track_entry.extend_from_slice(&entry);
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::TRACKS).unwrap();
ebml::write_size(&mut out, track_entry.len() as u64).unwrap();
out.extend_from_slice(&track_entry);
let e = match MkvStream::open(Cursor::new(out)) {
Ok(_) => panic!("expected MkvInvalid, got Ok"),
Err(e) => e,
};
assert!(is_mkv_invalid(&e));
}
#[test]
fn block_group_frame_round_trips_with_duration() {
let block = [0x82u8, 0x00, 0x05, 0x00, 0x11, 0x22, 0x33]; let mut bg_body = Vec::new();
ebml::write_id(&mut bg_body, ebml::BLOCK).unwrap();
ebml::write_size(&mut bg_body, block.len() as u64).unwrap();
bg_body.extend_from_slice(&block);
ebml::write_uint(&mut bg_body, ebml::BLOCK_DURATION, 40).unwrap();
let mut cluster = Vec::new();
ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap();
ebml::write_unknown_size(&mut cluster).unwrap();
ebml::write_uint(&mut cluster, ebml::CLUSTER_TIMESTAMP, 100).unwrap();
ebml::write_id(&mut cluster, ebml::BLOCK_GROUP).unwrap();
ebml::write_size(&mut cluster, bg_body.len() as u64).unwrap();
cluster.extend_from_slice(&bg_body);
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, 0).unwrap();
let mut tracks = Vec::new();
for (n, t) in [(1u64, 1u64), (2u64, 2u64)] {
let mut entry = Vec::new();
ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, n).unwrap();
ebml::write_uint(&mut entry, ebml::TRACK_TYPE, t).unwrap();
ebml::write_id(&mut tracks, ebml::TRACK_ENTRY).unwrap();
ebml::write_size(&mut tracks, entry.len() as u64).unwrap();
tracks.extend_from_slice(&entry);
}
ebml::write_id(&mut out, ebml::TRACKS).unwrap();
ebml::write_size(&mut out, tracks.len() as u64).unwrap();
out.extend_from_slice(&tracks);
out.extend_from_slice(&cluster);
let mut stream = MkvStream::open(Cursor::new(out)).unwrap();
let frame = stream
.read()
.unwrap()
.expect("BlockGroup frame must be read");
assert_eq!(frame.track, 1, "track 2 → index 1");
assert!(!frame.keyframe);
assert_eq!(frame.data, vec![0x11, 0x22, 0x33]);
assert_eq!(frame.pts, 105 * 1_000_000, "pts = (cluster 100 + rel 5) ms");
assert_eq!(frame.duration_ns, Some(40 * 1_000_000));
}
#[test]
fn track_number_zero_is_rejected() {
let bytes = mkv_with_track_and_cluster(0, 1, &[]);
let e = open_err(MkvStream::open(Cursor::new(bytes)));
assert!(is_mkv_invalid(&e));
}
#[test]
fn track_number_above_u16_is_rejected() {
let bytes = mkv_with_track_and_cluster(65536, 1, &[]);
let e = open_err(MkvStream::open(Cursor::new(bytes)));
assert!(is_mkv_invalid(&e));
}
#[test]
fn unknown_size_inner_child_in_tracks_is_rejected() {
let mut entry = Vec::new();
ebml::write_id(&mut entry, ebml::TRACK_NUMBER).unwrap();
ebml::write_unknown_size(&mut entry).unwrap();
let mut tracks = Vec::new();
ebml::write_id(&mut tracks, ebml::TRACK_ENTRY).unwrap();
ebml::write_size(&mut tracks, entry.len() as u64).unwrap();
tracks.extend_from_slice(&entry);
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::TRACKS).unwrap();
ebml::write_size(&mut out, tracks.len() as u64).unwrap();
out.extend_from_slice(&tracks);
let e = open_err(MkvStream::open(Cursor::new(out)));
assert!(is_mkv_invalid(&e));
}
#[test]
fn oversized_title_string_is_rejected() {
let mut info = Vec::new();
ebml::write_id(&mut info, ebml::TITLE).unwrap();
ebml::write_size(&mut info, MAX_STRING_LEN + 1).unwrap();
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, info.len() as u64).unwrap();
out.extend_from_slice(&info);
let e = match MkvStream::open(Cursor::new(out)) {
Ok(_) => panic!("expected MkvInvalid, got Ok"),
Err(e) => e,
};
assert!(is_mkv_invalid(&e));
}
#[test]
fn read_uint_val_len_nine_errors_not_panics() {
let mut data = Cursor::new(vec![0u8; 16]);
let e = ebml::read_uint_val(&mut data, 9).unwrap_err();
assert!(is_mkv_invalid(&e));
}
#[test]
fn read_float_val_bad_width_errors() {
let mut data = Cursor::new(vec![0u8; 16]);
let e = ebml::read_float_val(&mut data, 5).unwrap_err();
assert!(is_mkv_invalid(&e));
let mut z = Cursor::new(vec![0u8; 16]);
assert_eq!(ebml::read_float_val(&mut z, 0).unwrap(), 0.0);
let mut f4 = Cursor::new(vec![0u8; 16]);
assert!(ebml::read_float_val(&mut f4, 4).is_ok());
let mut f8 = Cursor::new(vec![0u8; 16]);
assert!(ebml::read_float_val(&mut f8, 8).is_ok());
}
#[test]
fn non_utf8_string_element_is_rejected() {
let mut data = Cursor::new(vec![0xFF, 0xFE, 0xFD, 0xFC]);
let e = ebml::read_string_val(&mut data, 4).unwrap_err();
assert!(is_mkv_invalid(&e));
}
#[test]
fn simple_block_track_zero_is_skipped() {
let mut cluster = Vec::new();
ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap();
ebml::write_unknown_size(&mut cluster).unwrap();
let bad = [0x80u8, 0x00, 0x00, 0x80, 0xEE];
ebml::write_id(&mut cluster, ebml::SIMPLE_BLOCK).unwrap();
ebml::write_size(&mut cluster, bad.len() as u64).unwrap();
cluster.extend_from_slice(&bad);
let good = [0x81u8, 0x00, 0x00, 0x80, 0xAB, 0xCD];
ebml::write_id(&mut cluster, ebml::SIMPLE_BLOCK).unwrap();
ebml::write_size(&mut cluster, good.len() as u64).unwrap();
cluster.extend_from_slice(&good);
let bytes = mkv_with_track_and_cluster(1, 1, &cluster);
let mut stream = MkvStream::open(Cursor::new(bytes)).unwrap();
let frame = stream.read().unwrap().expect("track-1 frame expected");
assert_eq!(frame.track, 0);
assert_eq!(frame.data, vec![0xAB, 0xCD]);
}
#[test]
fn block_vint_width_selection_and_values() {
assert_eq!(block_vint(&[0x81]), (1, 1));
assert_eq!(block_vint(&[0xFF]), (0x7F, 1)); assert_eq!(block_vint(&[0x40, 0x80]), (0x80, 2));
assert_eq!(block_vint(&[0x7F, 0xFF]), (0x3FFF, 2)); assert_eq!(block_vint(&[0x20, 0x00, 0x01]), (1, 3));
assert_eq!(block_vint(&[0x3F, 0xFF, 0xFF]), (0x1F_FFFF, 3));
assert_eq!(block_vint(&[0x10, 0x00, 0x00, 0x01]), (1, 4));
assert_eq!(block_vint(&[0x1F, 0xFF, 0xFF, 0xFF]), (0x0FFF_FFFF, 4));
}
#[test]
fn block_vint_unsupported_and_truncated_forms() {
assert_eq!(block_vint(&[]), (0, 0));
assert_eq!(block_vint(&[0x40]), (0, 1));
assert_eq!(block_vint(&[0x08, 0, 0, 0, 0]), (0, 1));
assert_eq!(block_vint(&[0x00, 0x11]), (0, 1));
}
#[test]
fn parse_block_too_short_is_none() {
assert!(parse_block(&[0x81, 0x00, 0x00], 0, 1_000_000, 1, None).is_none());
assert!(parse_block(&[], 0, 1_000_000, 1, None).is_none());
}
#[test]
fn parse_block_header_longer_than_payload_is_none() {
let block = [0x40u8, 0x01, 0x00, 0x00]; assert!(parse_block(&block, 0, 1_000_000, 2, None).is_none());
}
#[test]
fn parse_block_track_index_out_of_range_is_none() {
let block = [0x82u8, 0x00, 0x00, 0x80, 0xAA]; assert!(parse_block(&block, 0, 1_000_000, 1, None).is_none());
let f = parse_block(&block, 0, 1_000_000, 2, None).unwrap();
assert_eq!(f.track, 1);
}
#[test]
fn parse_block_pts_honours_timestamp_scale() {
let block = [0x81u8, 0x00, 0x0A, 0x80, 0xAA]; let f = parse_block(&block, 100, 1_000_000, 1, None).unwrap();
assert_eq!(f.pts, 110 * 1_000_000);
assert!(f.keyframe);
let f = parse_block(&block, 100, 90_000, 1, None).unwrap();
assert_eq!(f.pts, 110 * 90_000);
}
#[test]
fn parse_block_negative_rel_ts_is_signed() {
let block = [0x81u8, 0xFF, 0xFF, 0x80, 0xAA]; let f = parse_block(&block, 100, 1_000_000, 1, None).unwrap();
assert_eq!(f.pts, 99 * 1_000_000, "rel_ts -1 must subtract one tick");
}
#[test]
fn parse_block_keyframe_flag_and_duration_propagate() {
let kf = [0x81u8, 0x00, 0x00, 0x80, 0xAA];
let nkf = [0x81u8, 0x00, 0x00, 0x00, 0xAA];
assert!(parse_block(&kf, 0, 1_000_000, 1, None).unwrap().keyframe);
assert!(!parse_block(&nkf, 0, 1_000_000, 1, None).unwrap().keyframe);
let f = parse_block(&kf, 0, 1_000_000, 1, Some(40_000_000)).unwrap();
assert_eq!(f.duration_ns, Some(40_000_000));
}
#[test]
fn parse_block_pts_saturates_no_overflow() {
let block = [0x81u8, 0x00, 0x00, 0x80, 0xAA];
let f = parse_block(&block, i64::MAX, 1_000_000, 1, None).unwrap();
assert_eq!(f.pts, i64::MAX, "ticks→ns must saturate, not wrap/panic");
}
#[test]
fn cluster_timestamp_above_i64_max_is_rejected() {
let mut cluster = Vec::new();
ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap();
ebml::write_unknown_size(&mut cluster).unwrap();
ebml::write_id(&mut cluster, ebml::CLUSTER_TIMESTAMP).unwrap();
ebml::write_size(&mut cluster, 8).unwrap();
cluster.extend_from_slice(&0xFFFF_FFFF_FFFF_FFFFu64.to_be_bytes());
let bytes = mkv_with_track_and_cluster(1, 1, &cluster);
let mut stream = MkvStream::open(Cursor::new(bytes)).unwrap();
let e = stream.read().unwrap_err();
assert!(is_mkv_invalid(&e));
}
#[test]
fn zero_timestamp_scale_clamps_to_default() {
let mut info = Vec::new();
ebml::write_uint(&mut info, ebml::TIMESTAMP_SCALE, 0).unwrap();
let mut entry = Vec::new();
ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, 1).unwrap();
ebml::write_uint(&mut entry, ebml::TRACK_TYPE, 1).unwrap();
let mut track_entry = Vec::new();
ebml::write_id(&mut track_entry, ebml::TRACK_ENTRY).unwrap();
ebml::write_size(&mut track_entry, entry.len() as u64).unwrap();
track_entry.extend_from_slice(&entry);
let mut cluster = Vec::new();
ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap();
ebml::write_unknown_size(&mut cluster).unwrap();
ebml::write_uint(&mut cluster, ebml::CLUSTER_TIMESTAMP, 100).unwrap();
let block = [0x81u8, 0x00, 0x05, 0x80, 0xAA];
ebml::write_id(&mut cluster, ebml::SIMPLE_BLOCK).unwrap();
ebml::write_size(&mut cluster, block.len() as u64).unwrap();
cluster.extend_from_slice(&block);
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, info.len() as u64).unwrap();
out.extend_from_slice(&info);
ebml::write_id(&mut out, ebml::TRACKS).unwrap();
ebml::write_size(&mut out, track_entry.len() as u64).unwrap();
out.extend_from_slice(&track_entry);
out.extend_from_slice(&cluster);
let mut stream = MkvStream::open(Cursor::new(out)).unwrap();
let f = stream.read().unwrap().expect("frame");
assert_eq!(f.pts, 105 * 1_000_000, "zero scale must clamp to 1ms");
}
#[test]
fn duration_uses_timestamp_scale_for_seconds() {
let mut info = Vec::new();
ebml::write_uint(&mut info, ebml::TIMESTAMP_SCALE, 1_000_000).unwrap();
ebml::write_float(&mut info, ebml::DURATION, 5000.0).unwrap();
let mut entry = Vec::new();
ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, 1).unwrap();
ebml::write_uint(&mut entry, ebml::TRACK_TYPE, 1).unwrap();
let mut track_entry = Vec::new();
ebml::write_id(&mut track_entry, ebml::TRACK_ENTRY).unwrap();
ebml::write_size(&mut track_entry, entry.len() as u64).unwrap();
track_entry.extend_from_slice(&entry);
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, info.len() as u64).unwrap();
out.extend_from_slice(&info);
ebml::write_id(&mut out, ebml::TRACKS).unwrap();
ebml::write_size(&mut out, track_entry.len() as u64).unwrap();
out.extend_from_slice(&track_entry);
let stream = MkvStream::open(Cursor::new(out)).unwrap();
assert_eq!(stream.info().duration_secs, 5.0);
}
#[test]
fn missing_ebml_header_is_rejected() {
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap(); ebml::write_size(&mut out, 0).unwrap();
let e = open_err(MkvStream::open(Cursor::new(out)));
assert!(is_mkv_invalid(&e));
}
#[test]
fn segment_must_follow_ebml_header() {
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap(); ebml::write_size(&mut out, 0).unwrap();
let e = open_err(MkvStream::open(Cursor::new(out)));
assert!(is_mkv_invalid(&e));
}
#[test]
fn track_type_to_codec_and_pid_mapping_round_trips() {
let mut entry = Vec::new();
ebml::write_uint(&mut entry, ebml::TRACK_NUMBER, 1).unwrap();
ebml::write_uint(&mut entry, ebml::TRACK_TYPE, 1).unwrap();
ebml::write_string(&mut entry, ebml::CODEC_ID, "V_MPEGH/ISO/HEVC").unwrap();
let mut track_entry = Vec::new();
ebml::write_id(&mut track_entry, ebml::TRACK_ENTRY).unwrap();
ebml::write_size(&mut track_entry, entry.len() as u64).unwrap();
track_entry.extend_from_slice(&entry);
let mut out = Vec::new();
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, 0).unwrap();
ebml::write_id(&mut out, ebml::TRACKS).unwrap();
ebml::write_size(&mut out, track_entry.len() as u64).unwrap();
out.extend_from_slice(&track_entry);
let stream = MkvStream::open(Cursor::new(out)).unwrap();
match &stream.info().streams[0] {
crate::disc::Stream::Video(v) => {
assert_eq!(v.codec, Codec::Hevc);
assert_eq!(v.pid, 0x1011);
}
_ => panic!("expected video stream"),
}
}
#[test]
fn block_group_unknown_size_is_rejected() {
let mut cluster = Vec::new();
ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap();
ebml::write_unknown_size(&mut cluster).unwrap();
ebml::write_id(&mut cluster, ebml::BLOCK_GROUP).unwrap();
ebml::write_unknown_size(&mut cluster).unwrap(); let bytes = mkv_with_track_and_cluster(1, 1, &cluster);
let mut stream = MkvStream::open(Cursor::new(bytes)).unwrap();
let e = stream.read().unwrap_err();
assert!(is_mkv_invalid(&e));
}
#[test]
fn read_then_eof_returns_none() {
let mut cluster = Vec::new();
ebml::write_id(&mut cluster, ebml::CLUSTER).unwrap();
ebml::write_unknown_size(&mut cluster).unwrap();
let block = [0x81u8, 0x00, 0x00, 0x80, 0xAA];
ebml::write_id(&mut cluster, ebml::SIMPLE_BLOCK).unwrap();
ebml::write_size(&mut cluster, block.len() as u64).unwrap();
cluster.extend_from_slice(&block);
let bytes = mkv_with_track_and_cluster(1, 1, &cluster);
let mut stream = MkvStream::open(Cursor::new(bytes)).unwrap();
assert!(stream.read().unwrap().is_some(), "first frame");
assert!(stream.read().unwrap().is_none(), "clean EOF → None");
}
}