pq/libfreemkv
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

libfreemkv: v1.0 hardening — codec/EBML/TS robustness + DTS parser fixes

Browse Source 
Audit-driven fixes (rounds 1–3):
- hevc: correct hvcC profile/level SPS offsets (HEVC has a 2-byte NAL header)
- mkv: map all DTS variants to the registered A_DTS codec id; force a new
  cluster before the i16 cluster-relative timestamp can overflow
- ebml/mkvstream: bound untrusted EBML sizes (no multi-GB allocs); reject
  uint>8 (was an OOB panic) and non-{0,4,8} float widths (were a desync)
- ts: skip PES-header bytes that span a TS packet boundary; add the PMT
  section_len/prog_info_len bounds the PAT parser already had
- ac3: preserve a 0x0B77 syncword split across a PES boundary; cap buffer
- dts: validate each next-core boundary by decoded core size (a 0x7FFE8001
  pattern inside XLL payload no longer false-splits/drops the lossless
  extension); reject sub-minimum core frames; fix forced-emit PTS base
- lpcm: DVD program-stream PCM no longer double-strips the BD LPCM header
- vc1/mpeg2: do not emit a parameter-set-only PES as a standalone frame
- pgs/truehd: cap the pending reassembly buffer (parity with ac3/dts)
- aacs: ts_syncs_intact uses the exact packet count
- prefetched: capacity-guard the recycled-buffer set_len
- Cargo.toml: exclude project docs from the published crate

Convergence: a third independent audit pass found no remaining material
(CRITICAL/HIGH/MEDIUM) issues. Full precommit (fmt + clippy -D + tests,
Rust 1.86) green.
This commit is contained in:
MattJackson
2026-06-05 16:23:39 -07:00
parent e2aa9abd6d
commit 6be5198886
18 changed files with 1201 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.toml
+2
View File
@@ -8,6 +8,8 @@ description = "Open source raw disc access library for optical drives"
repository = "https://github.com/freemkv/libfreemkv"
keywords = ["bluray", "uhd", "optical", "scsi", "disc"]
categories = ["hardware-support", "multimedia"]
# Keep internal AI-instruction / private notes out of the published crate.
exclude = []
[dependencies]
serde = { version = "1", features = ["derive"] }
src/aacs/decrypt.rs
+6 -1
View File
@@ -100,7 +100,12 @@ fn ts_syncs_intact(unit: &[u8]) -> bool {
}
offset += TS_PACKET_LEN;
}
let total = (unit.len() - 4) / TS_PACKET_LEN + 1;
// One sync byte is checked per 192-byte BD-TS packet (at offset 4 of
// each). `total` is exactly that packet count; the old
// `(len - 4) / TS_PACKET_LEN + 1` over-counted by one for lengths of
// the form `4 + k·192` (harmless for the always-6144 aligned unit, but
// wrong in general and it biased the majority threshold).
let total = unit.len() / TS_PACKET_LEN;
count > total / 2
}
src/mux/codec/ac3.rs
+119 -7
View File
@@ -6,6 +6,14 @@
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
/// Hard cap on the carry-over buffer. An AC-3/E-AC-3 frame is at most 8192
/// bytes (the `frame_size > 8192` reject below), so a single straddling frame
/// plus a little slack never needs more than this. If the buffer grows past
/// the cap without yielding a frame (pathological / never-syncing input) we
/// drop it and resync rather than accumulate one PES worth of data per call
/// for the whole title.
const MAX_AC3_BUF: usize = 64 * 1024;
pub struct Ac3Parser {
/// Leftover bytes from previous PES (incomplete frame at end).
buf: Vec<u8>,
@@ -81,19 +89,38 @@ impl CodecParser for Ac3Parser {
pos = start + frame_size;
}
// Keep unconsumed data for next call
// `pos` points to the start of unconsumed data (either a partial sync or leftover)
// Keep unconsumed data for the next call. `pos` is the start of the
// unconsumed region: either a partial frame that straddles this PES
// boundary — which, by construction, begins at a syncword (every byte
// before `pos` was emitted as a frame or skipped as pre-sync junk) — or
// trailing bytes too short to size/complete a frame. Carry from `pos`,
// NOT from the next syncword: discarding bytes between `pos` and the
// next sync would drop the partial frame we are deliberately keeping
// across the boundary.
let keep_from = if pos < data.len() {
// Find the last sync word position in the unconsumed region
find_ac3_sync(&data[pos..])
.map(|o| pos + o)
.unwrap_or(data.len())
// A syncword at/after `pos` marks the carry-over start (anything
// before it is junk with no sync). With no full sync, retain the
// whole tail — including a lone trailing 0x0B that may be the first
// half of a syncword split across the PES boundary.
match find_ac3_sync(&data[pos..]) {
Some(o) => pos + o,
None if data.last() == Some(&0x0B) => data.len() - 1,
None => data.len(),
}
} else {
data.len()
};
if keep_from < data.len() {
self.buf = data[keep_from..].to_vec();
let tail = &data[keep_from..];
if tail.len() > MAX_AC3_BUF {
// No frame could be parsed out of a buffer this large — this is
// not valid AC-3 here. Drop it and resync on the next PES rather
// than grow without bound on pathological input.
self.buf.clear();
} else {
self.buf = tail.to_vec();
}
} else {
self.buf.clear();
}
@@ -276,6 +303,91 @@ mod tests {
assert_eq!(frames[0].data.len(), 160);
}
#[test]
fn sync_word_split_across_pes_is_preserved() {
// A frame whose 0x0B77 syncword straddles the PES boundary (0x0B at the
// tail of PES 1, 0x77 at the head of PES 2) must still be emitted whole.
// Previously the lone trailing 0x0B was dropped and the frame lost.
let mut parser = Ac3Parser::new();
let frame_data = make_ac3_frame(0, 2); // 160 bytes, starts with 0x0B 0x77
// PES 1: a complete frame, then a single 0x0B (first half of next sync).
let mut pes1_data = frame_data.clone();
pes1_data.push(0x0B);
let pes1 = PesPacket {
pid: 0,
pts: Some(90000),
dts: None,
data: pes1_data,
};
let frames1 = parser.parse(&pes1);
assert_eq!(frames1.len(), 1, "first complete frame emitted");
// PES 2: 0x77 (second half of sync) + rest of the second frame.
let mut pes2_data = vec![0x77];
pes2_data.extend_from_slice(&frame_data[2..]);
let pes2 = PesPacket {
pid: 0,
pts: Some(93000),
dts: None,
data: pes2_data,
};
let frames2 = parser.parse(&pes2);
assert_eq!(frames2.len(), 1, "split-sync frame must be recovered");
assert_eq!(frames2[0].data.len(), 160);
}
#[test]
fn buffer_stays_bounded_across_many_garbage_pes() {
// Finding 14: the carry-over buffer must never grow without bound. Feed
// many large PES packets that contain no usable frame and assert the
// retained buffer stays tiny — carry-from-`pos` drops all pre-sync junk,
// and a never-completing frame is bounded by the 8192-byte frame cap and
// the MAX_AC3_BUF resync guard.
let mut parser = Ac3Parser::new();
for i in 0..256 {
// Vary the trailing byte so we also exercise the lone-0x0B retain.
let mut data = vec![0x55u8; 8192];
if i % 3 == 0 {
*data.last_mut().unwrap() = 0x0B;
}
let pes = PesPacket {
pid: 0,
pts: None,
dts: None,
data,
};
let frames = parser.parse(&pes);
assert!(frames.is_empty());
assert!(
parser.buf.len() <= MAX_AC3_BUF,
"buffer grew to {} (cap {})",
parser.buf.len(),
MAX_AC3_BUF
);
}
// After all that garbage the retained tail is at most a single partial
// syncword byte — never an accumulation of whole PES packets.
assert!(parser.buf.len() <= 1, "retained {} bytes", parser.buf.len());
}
#[test]
fn split_sync_below_cap_is_still_retained() {
// The cap must not break the normal split-sync straddle: a short tail
// ending in 0x0B (well under the cap) is retained so the next PES can
// complete the syncword.
let mut parser = Ac3Parser::new();
let data = vec![0x00, 0x00, 0x0B];
let pes = PesPacket {
pid: 0,
pts: None,
dts: None,
data,
};
assert!(parser.parse(&pes).is_empty());
assert_eq!(parser.buf, vec![0x0B], "lone trailing 0x0B retained");
}
#[test]
fn ac3_frame_size_table() {
// fscod=0 (48kHz), frmsizecod=0: 64 words = 128 bytes
src/mux/codec/dts.rs
+201 -14
View File
@@ -44,6 +44,19 @@ impl DtsParser {
/// this without a clean boundary we resync rather than stall or balloon.
const MAX_AU_BYTES: usize = 65536;
/// Minimum plausible DTS core frame size. The core header alone is ~10-14
/// bytes; a decoded `core_size` below this means we matched a false/corrupt
/// core sync (the 14-bit `fsize` field decoded to a tiny value) rather than a
/// real frame, so we resync instead of emitting a junk access unit.
const MIN_CORE_FRAME_BYTES: usize = 10;
/// Sentinel for "no valid PTS base captured yet". Real PTS-in-ns values are
/// non-negative (derived from the unsigned 90 kHz PES timestamp), so a negative
/// value can never collide with a genuine timestamp. Used to mark the PTS base
/// invalid after a forced flush so the next PES sets it regardless of buffer
/// state.
const PTS_UNSET: i64 = -1;
impl CodecParser for DtsParser {
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
if pes.data.is_empty() {
@@ -65,7 +78,11 @@ impl CodecParser for DtsParser {
// old per-PES emit that dropped the extension PES packets and
// downgraded the track to lossy DTS core (the Dunkirk / Fight Club
// bug). The PTS is the core frame's PTS, captured when the unit began.
if self.buf.is_empty() {
// Capture the access unit's PTS base on a fresh buffer, or whenever a
// prior forced (safety-valve) flush left it invalidated — in the
// forced case the bytes still in `buf` are not a real core frame, so
// the first PES to arrive after the flush carries the correct base.
if self.buf.is_empty() || self.pending_pts == PTS_UNSET {
self.pending_pts = pts_ns;
}
self.buf.extend_from_slice(&pes.data);
@@ -96,7 +113,15 @@ impl CodecParser for DtsParser {
break;
}
let core_size = dts_core_frame_size(&self.buf);
if core_size == 0 || core_size > MAX_AU_BYTES {
// `dts_core_frame_size` returns a 14-bit `fsize + 1`, so it is
// always in [1, 16384]; the bare `== 0` / `> MAX_AU_BYTES` checks
// can never fire. A real DTS core header is at least ~10-14 bytes,
// so any decoded size below that came from a false/corrupt sync.
// Reject it (drain the 4 syncword bytes and resync) instead of
// letting a tiny bogus size close the current access unit at a junk
// boundary and drop the trailing extension substreams. The
// `> MAX_AU_BYTES` upper bound is kept as a harmless guard.
if !(MIN_CORE_FRAME_BYTES..=MAX_AU_BYTES).contains(&core_size) {
// Bogus core sync — skip past it and resync.
self.buf.drain(..4);
continue;
@@ -105,12 +130,25 @@ impl CodecParser for DtsParser {
break; // core frame not fully buffered yet — wait
}
// The access unit ends at the next core sync. Search begins after
// this core's syncword so we don't re-match it. Anything between
// the core and that next sync is this unit's extension substream(s).
let au_end = match find_sync(&self.buf[core_size..], &DTS_CORE_SYNC) {
Some(rel) => core_size + rel,
None => {
// The access unit ends at the next *valid* core sync. The search
// begins after this core's syncword so we don't re-match it.
// Anything between the core and that next sync is this unit's
// extension substream(s) — which can themselves contain byte
// sequences matching the core syncword, so a raw `find_sync` match
// is not enough: a candidate is only a real boundary if its decoded
// core size is plausible. `next_core_boundary` skips bogus matches.
//
// `forced` distinguishes a real next-core boundary from a forced
// safety-valve flush. On a forced flush the access unit was NOT
// closed by a new core sync, so the bytes following it are not a
// fresh core frame and the current PES's PTS (which on a forced
// flush is an extension-substream PES, carrying its own later
// timestamp) must NOT become the next unit's PTS base.
let mut forced = false;
let au_end = match next_core_boundary(&self.buf, core_size) {
NextCore::Found(end) => end,
NextCore::NeedMore => break, // candidate sync needs more header
NextCore::None => {
// No next core sync buffered yet. The trailing extension
// substream PES packets may still be arriving, so WAIT for
// them rather than emit a core-only (lossy) frame — unless
@@ -119,6 +157,7 @@ impl CodecParser for DtsParser {
if self.buf.len() <= MAX_AU_BYTES {
break;
}
forced = true;
self.buf.len()
}
};
@@ -131,12 +170,19 @@ impl CodecParser for DtsParser {
duration_ns: None,
});
self.buf.drain(..au_end);
// The next access unit (now at buf start) belongs to a later PTS.
// We can't know it exactly until its core PES arrives, but the
// current PES's PTS is the best available approximation when the
// boundary fell inside this PES; refine on the next call's start.
if forced {
// Safety-valve flush: the next access unit's real core PES has
// not arrived. Invalidate the PTS so the next PES sets it
// regardless of buffer state, rather than inheriting this
// (non-core) PES's timestamp.
self.pending_pts = PTS_UNSET;
} else {
// Real boundary: the next access unit (now at buf start) begins
// at a core sync that arrived inside this PES, so this PES's PTS
// is the correct base for it.
self.pending_pts = pts_ns;
}
}
frames
}
@@ -151,12 +197,21 @@ impl CodecParser for DtsParser {
return Vec::new();
}
let core_size = dts_core_frame_size(&self.buf);
if core_size == 0 || self.buf.len() < core_size {
// `dts_core_frame_size` returns a 14-bit `fsize + 1` (never 0), so the
// old `== 0` check was dead; reject a sub-minimum core like `parse()`.
if core_size < MIN_CORE_FRAME_BYTES || self.buf.len() < core_size {
self.buf.clear();
return Vec::new();
}
let au = std::mem::take(&mut self.buf);
let pts_ns = self.pending_pts;
// A non-empty buffer here means a PES arrived after any prior forced
// flush (which fully drains `buf`), so `pending_pts` was reset to that
// PES's real PTS. Clamp the sentinel to 0 defensively all the same.
let pts_ns = if self.pending_pts == PTS_UNSET {
0
} else {
self.pending_pts
};
vec![Frame {
pts_ns,
keyframe: true,
@@ -177,6 +232,40 @@ fn find_sync(data: &[u8], pattern: &[u8; 4]) -> Option<usize> {
(0..=data.len() - 4).find(|&i| data[i..i + 4] == *pattern)
}
/// Result of scanning for the next valid core sync that closes an access unit.
enum NextCore {
/// A valid next core sync was found; the access unit ends at this offset.
Found(usize),
/// A candidate core sync was found but its header isn't fully buffered yet,
/// so its validity can't be decided — wait for more data.
NeedMore,
/// No (further) core sync found in the buffer.
None,
}
/// Find the next *valid* core sync after the current core frame, to delimit the
/// access unit. Extension-substream payload can contain byte sequences that
/// match the core syncword, so each candidate is validated by decoding its
/// core size: a match whose decoded size is implausible (< MIN_CORE_FRAME_BYTES
/// or > MAX_AU_BYTES) is a false sync and is skipped, continuing the search.
fn next_core_boundary(buf: &[u8], core_size: usize) -> NextCore {
let mut from = core_size;
while let Some(rel) = find_sync(&buf[from..], &DTS_CORE_SYNC) {
let pos = from + rel;
// Need the candidate's core header to judge it.
if buf.len() - pos < 10 {
return NextCore::NeedMore;
}
let sz = dts_core_frame_size(&buf[pos..]);
if (MIN_CORE_FRAME_BYTES..=MAX_AU_BYTES).contains(&sz) {
return NextCore::Found(pos);
}
// False sync inside extension payload — skip it and keep searching.
from = pos + 4;
}
NextCore::None
}
/// DTS core frame size from header bits.
/// fsize is at bits 46-59 (14 bits) of the header: bytes 5-7.
fn dts_core_frame_size(data: &[u8]) -> usize {
@@ -349,6 +438,104 @@ mod tests {
assert_eq!(tail[0].data.len(), 512 + 300);
}
/// Build 4 bytes that look like a DTS core sync but whose `fsize` field
/// decodes to a tiny `core_size` (< MIN_CORE_FRAME_BYTES). With the
/// dead-code guards this passed validation and could close an access unit
/// at a junk boundary; with the fix it must be drained and resynced past.
fn bogus_tiny_core_sync() -> Vec<u8> {
// Core sync + zero header bytes. fsize = 0 → core_size = 1 (< 10).
let mut v = vec![0u8; 10];
v[0..4].copy_from_slice(&DTS_CORE_SYNC);
// bytes 5,6,7 left zero → fsize = 0 → dts_core_frame_size = 1.
assert_eq!(dts_core_frame_size(&v), 1);
v
}
#[test]
fn bogus_tiny_core_sync_does_not_split_or_drop_real_au() {
// A real core frame followed by an extension substream that happens to
// contain a false core sync whose fsize decodes tiny. The bogus sync
// must NOT close the real access unit early (dropping the rest of the
// extension) nor emit a junk few-byte frame — it must be skipped, and
// the whole core + extension preserved as one access unit.
let mut parser = DtsParser::new();
// Frame 1: core(512) + an extension whose body embeds a bogus tiny
// core sync midway through.
let mut ext = make_dts_ext(256);
// Embed the bogus core sync inside the extension body (offset 64).
let bogus = bogus_tiny_core_sync();
ext[64..64 + bogus.len()].copy_from_slice(&bogus);
let mut frame1 = make_dts_core(512);
frame1.extend_from_slice(&ext);
// No next REAL core yet → frame 1 held.
assert!(
parser.parse(&make_pes(frame1, Some(90000))).is_empty(),
"bogus tiny core sync must not close the AU; wait for a real core"
);
// Frame 2's real core arrives — closes frame 1.
let f = parser.parse(&make_pes(make_dts_core(640), Some(93000)));
assert_eq!(f.len(), 1, "exactly one real access unit emitted");
assert_eq!(
f[0].data.len(),
512 + 256,
"AU must be the full core + extension, not split at the bogus sync"
);
assert_eq!(f[0].pts_ns, pts_to_ns(90000), "AU keeps the core's PTS");
let tail = parser.flush();
assert_eq!(tail.len(), 1);
assert_eq!(tail[0].data.len(), 640);
}
#[test]
fn forced_emit_does_not_corrupt_next_au_pts() {
// When the buffer exceeds MAX_AU_BYTES with no next core sync, the
// parser force-emits for forward progress. The current PES at that
// point is an extension-substream PES (later PTS). The forced path must
// NOT make that extension PTS the base of the NEXT access unit.
let mut parser = DtsParser::new();
// Core PES at the real PTS, then a giant extension (no next core) that
// pushes the buffer past MAX_AU_BYTES, forcing an emit.
let core_pts = 90000i64;
assert!(
parser
.parse(&make_pes(make_dts_core(512), Some(core_pts)))
.is_empty()
);
let ext_pts = 120000i64; // later extension-PES timestamp
let big_ext = make_dts_ext(MAX_AU_BYTES + 1024);
let f = parser.parse(&make_pes(big_ext, Some(ext_pts)));
assert_eq!(f.len(), 1, "oversized buffer force-emits one AU");
assert_eq!(
f[0].pts_ns,
pts_to_ns(core_pts),
"forced AU keeps the core PTS"
);
// The next REAL core PES arrives with its own PTS. Its AU must inherit
// THIS core's PTS, not the prior extension PES timestamp.
let next_core_pts = 150000i64;
assert!(
parser
.parse(&make_pes(make_dts_core(512), Some(next_core_pts)))
.is_empty()
);
let next_next_pts = 180000i64;
let f2 = parser.parse(&make_pes(make_dts_core(512), Some(next_next_pts)));
assert_eq!(f2.len(), 1);
assert_eq!(
f2[0].pts_ns,
pts_to_ns(next_core_pts),
"AU after a forced emit must use the next core's PTS, not the \
stale extension PTS"
);
}
#[test]
fn codec_private_none() {
let parser = DtsParser::new();
src/mux/codec/hevc.rs
+131 -22
View File
@@ -62,6 +62,11 @@ impl CodecParser for HevcParser {
while let Some(sc_pos) = find_start_code(data, pos) {
if let Some(nal_start) = skip_start_code(data, sc_pos) {
let next = find_start_code(data, nal_start).unwrap_or(data.len());
// Strip the leading zeros of the following start code. For a
// conforming bitstream this is lossless: rbsp_trailing_bits()
// sets a stop-one bit, so the final byte of any RBSP is never
// 0x00 — the only trailing zeros here belong to the next
// 00 00 (00) 01 prefix.
let mut end = next;
while end > nal_start && data[end - 1] == 0x00 {
end -= 1;
@@ -124,34 +129,41 @@ impl CodecParser for HevcParser {
// Full HEVCDecoderConfigurationRecord is complex — for now, concatenate
let mut record = Vec::new();
// Minimal HEVCDecoderConfigurationRecord header
// Minimal HEVCDecoderConfigurationRecord header.
//
// The stored SPS NAL is [2-byte HEVC NAL header][SPS RBSP...].
// The RBSP begins at sps[2]; profile_tier_level() begins one byte
// later, after sps_video_parameter_set_id u(4) +
// sps_max_sub_layers_minus1 u(3) + sps_temporal_id_nesting_flag u(1)
// (= sps[2], a full byte). So the profile_tier_level fields are:
// sps[3] general_profile_space u(2)+tier u(1)+profile_idc u(5)
// sps[4..8] general_profile_compatibility_flags u(32)
// sps[8..14] general_constraint_indicator_flags 48 bits
// sps[14] general_level_idc u(8)
// (Byte-aligned read; emulation-prevention bytes within the first
// 15 SPS bytes are not handled — extremely rare and matches the
// pre-existing simplification.)
record.push(1); // configurationVersion
// General profile space, tier flag, profile IDC from SPS
if sps.len() > 3 {
record.push(sps[1]); // general_profile_space + general_tier_flag + general_profile_idc
// general_profile_space + general_tier_flag + general_profile_idc
record.push(if sps.len() > 3 { sps[3] } else { 0 });
// general_profile_compatibility_flags (4 bytes) — SPS bytes 4..8
if sps.len() > 7 {
record.extend_from_slice(&sps[4..8]);
} else {
record.push(0);
let avail = sps.len().saturating_sub(4).min(4);
record.extend_from_slice(&sps[sps.len().min(4)..sps.len().min(8)]);
record.extend_from_slice(&vec![0u8; 4 - avail]);
}
// general_profile_compatibility_flags (4 bytes) — from SPS bytes 2..6
if sps.len() > 5 {
record.extend_from_slice(&sps[2..6]);
// general_constraint_indicator_flags (6 bytes) — SPS bytes 8..14
if sps.len() > 13 {
record.extend_from_slice(&sps[8..14]);
} else {
record.extend_from_slice(&[0, 0, 0, 0]);
}
// general_constraint_indicator_flags (6 bytes) — from SPS bytes 6..12
if sps.len() > 11 {
record.extend_from_slice(&sps[6..12]);
} else {
let avail = sps.len().saturating_sub(6).min(6);
if avail > 0 {
record.extend_from_slice(&sps[6..6 + avail]);
let avail = sps.len().saturating_sub(8).min(6);
record.extend_from_slice(&sps[sps.len().min(8)..sps.len().min(14)]);
record.extend_from_slice(&vec![0u8; 6 - avail]);
} else {
record.extend_from_slice(&[0, 0, 0, 0, 0, 0]);
}
}
// general_level_idc
record.push(if sps.len() > 12 { sps[12] } else { 0 });
// general_level_idc — SPS byte 14
record.push(if sps.len() > 14 { sps[14] } else { 0 });
// min_spatial_segmentation_idc (4 + 12 bits)
record.extend_from_slice(&[0xF0, 0x00]);
// parallelismType (6 + 2 bits)
@@ -268,6 +280,103 @@ mod tests {
);
}
#[test]
fn hvcc_profile_tier_level_offsets() {
// The hvcC fixed header must read profile_tier_level from the SPS
// RBSP, not from the NAL header. Stored SPS = [2-byte NAL header][RBSP].
// RBSP layout (byte-aligned):
// sps[2] sps_vps_id/max_sub_layers/temporal_nesting
// sps[3] general_profile_space+tier+profile_idc
// sps[4..8] general_profile_compatibility_flags
// sps[8..14] general_constraint_indicator_flags
// sps[14] general_level_idc
let mut parser = HevcParser::new();
// Distinct, recognizable values for each field.
let sps_rbsp: [u8; 13] = [
0xAB, // sps[2] (vps_id etc.) — must NOT leak into profile fields
0x21, // sps[3] profile byte: space=0, tier=0, profile_idc=1
0x60, 0x00, 0x00, 0x00, // sps[4..8] compat flags
0x90, 0x00, 0x00, 0x00, 0x00, 0x00, // sps[8..14] constraint flags
0x7B, // sps[14] level_idc = 123
];
let mut data = Vec::new();
// VPS
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(32));
data.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
// SPS — 2-byte header + the structured RBSP above
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(33));
data.extend_from_slice(&sps_rbsp);
// PPS
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(34));
data.extend_from_slice(&[0xDD, 0xEE]);
let pes = make_pes(data, Some(0));
parser.parse(&pes);
let cp = parser
.codec_private()
.expect("codec_private should be Some");
// record[0] = configurationVersion
assert_eq!(cp[0], 1, "configurationVersion");
// record[1] = general_profile_space+tier+profile_idc <- sps[3]
assert_eq!(
cp[1], 0x21,
"profile byte must come from SPS RBSP, not NAL hdr"
);
// record[2..6] = general_profile_compatibility_flags <- sps[4..8]
assert_eq!(&cp[2..6], &[0x60, 0x00, 0x00, 0x00], "compatibility flags");
// record[6..12] = general_constraint_indicator_flags <- sps[8..14]
assert_eq!(
&cp[6..12],
&[0x90, 0x00, 0x00, 0x00, 0x00, 0x00],
"constraint flags"
);
// record[12] = general_level_idc <- sps[14]
assert_eq!(cp[12], 0x7B, "level_idc must come from sps[14]");
}
#[test]
fn hvcc_short_sps_does_not_panic() {
// A truncated SPS must still produce a fixed header without panicking
// and zero-pad the missing profile/level bytes.
let mut parser = HevcParser::new();
let mut data = Vec::new();
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(32));
data.extend_from_slice(&[0xAA]);
// SPS with only 3 RBSP bytes (stored len = 5): forces every guard path
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(33));
data.extend_from_slice(&[0x11, 0x22, 0x33]);
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(34));
data.extend_from_slice(&[0xDD]);
let pes = make_pes(data, Some(0));
parser.parse(&pes);
let cp = parser
.codec_private()
.expect("codec_private should be Some");
// sps stored = [hdr0, hdr1, 0x11, 0x22, 0x33], len 5.
// profile byte = sps[3] = 0x22; everything past sps[4]=0x33 is absent.
assert_eq!(cp[0], 1);
assert_eq!(cp[1], 0x22, "profile byte = sps[3]");
// compat flags: only sps[4]=0x33 present, rest zero-padded.
assert_eq!(&cp[2..6], &[0x33, 0x00, 0x00, 0x00]);
// constraint flags: none present, all zero.
assert_eq!(&cp[6..12], &[0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
// level_idc: absent, zero.
assert_eq!(cp[12], 0x00);
}
#[test]
fn codec_private_none_before_params() {
let parser = HevcParser::new();
src/mux/codec/lpcm.rs
+88 -10
View File
@@ -1,24 +1,39 @@
//! BD/DVD LPCM (Linear PCM) audio parser.
//!
//! BD LPCM PES packets have a 4-byte header:
//! BD LPCM PES packets (TS stream type 0x80) carry a 4-byte header on the
//! elementary-stream payload:
//! Bytes 0-1: audio frame number
//! Byte 2: reserved
//! Byte 3: quantization (bits 7-6), sample rate (bits 5-4), channel assignment (bits 3-0)
//! This header is part of the ES payload, so the BD parser must strip it.
//!
//! DVD LPCM (private stream 1, sub-stream 0xA0-0xA7) has a 3-byte header.
//! DVD LPCM lives in private stream 1 (sub-stream 0xA0-0xA7). Its 7-byte
//! private sub-header (sub_id + frames + first-access-unit-ptr(2) + emphasis +
//! quant/freq + channels) is stripped by `PsDemuxer` while demuxing the
//! Program Stream. By the time a DVD LPCM `PesPacket` reaches this parser its
//! `data` is already raw PCM, so the parser must NOT strip any further bytes —
//! doing so drops one sample pair per PES and drifts the audio.
//!
//! The raw PCM data follows the header. No framing is needed — each PES
//! payload minus its header is one complete audio frame.
//! The two origins are distinguished by the `strip_header` flag: BD = strip,
//! DVD = leave intact. The raw PCM data is otherwise one complete audio frame
//! per PES; no framing is needed.
//!
//! For MKV: codec ID "A_PCM/INT/BIG" (BD) or "A_PCM/INT/LIT" (DVD).
//! All frames are keyframes (uncompressed audio).
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
/// BD LPCM header size in bytes.
/// BD LPCM header size in bytes (present on BD-TS LPCM, absent on DVD-PS LPCM
/// because `PsDemuxer` already stripped the private sub-header).
const BD_LPCM_HEADER_SIZE: usize = 4;
pub struct LpcmParser;
pub struct LpcmParser {
/// Whether to strip the 4-byte BD LPCM header from each PES payload.
///
/// `true` for BD-TS LPCM (header still present), `false` for DVD-PS LPCM
/// (header already removed by `PsDemuxer`).
strip_header: bool,
}
impl Default for LpcmParser {
fn default() -> Self {
@@ -27,23 +42,36 @@ impl Default for LpcmParser {
}
impl LpcmParser {
/// BD-TS LPCM parser: strips the 4-byte BD LPCM header from each PES.
pub fn new() -> Self {
Self
Self { strip_header: true }
}
/// DVD-PS LPCM parser: `PsDemuxer` already stripped the private sub-header,
/// so the payload is raw PCM and no further bytes are removed.
pub fn new_dvd() -> Self {
Self {
strip_header: false,
}
}
}
impl CodecParser for LpcmParser {
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
// Skip the BD LPCM header (4 bytes).
let offset = if self.strip_header {
BD_LPCM_HEADER_SIZE
} else {
0
};
// If the PES is too short to contain header + data, return nothing.
if pes.data.len() <= BD_LPCM_HEADER_SIZE {
if pes.data.len() <= offset {
return Vec::new();
}
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
vec![Frame {
pts_ns,
keyframe: true,
data: pes.data[BD_LPCM_HEADER_SIZE..].to_vec(),
data: pes.data[offset..].to_vec(),
duration_ns: None,
}]
}
@@ -84,6 +112,56 @@ mod tests {
assert_eq!(frames[0].pts_ns, 1_000_000_000); // 90000 ticks = 1 second
}
#[test]
fn bd_lpcm_strips_4_byte_header() {
// BD-TS LPCM: the 4-byte BD header is part of the ES payload and must
// be stripped, leaving exactly the PCM bytes.
let mut parser = LpcmParser::new();
let header = vec![0x00, 0x01, 0x00, 0b1001_0001];
let pcm = vec![0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88];
let mut data = header;
data.extend_from_slice(&pcm);
let frames = parser.parse(&make_pes(data, Some(0)));
assert_eq!(frames.len(), 1);
assert_eq!(frames[0].data, pcm, "BD must strip exactly 4 header bytes");
}
#[test]
fn dvd_lpcm_preserves_all_pcm_bytes() {
// DVD-PS LPCM: PsDemuxer already removed the 7-byte private sub-header,
// so the payload handed to this parser is raw PCM. The DVD parser must
// NOT strip any further bytes (the round-2 audit Finding 3 bug: the BD
// 4-byte strip dropped one sample pair per PES, drifting the audio).
let mut parser = LpcmParser::new_dvd();
let pcm = vec![0xDE, 0xAD, 0xBE, 0xEF, 0xCA, 0xFE, 0x01, 0x02];
let frames = parser.parse(&make_pes(pcm.clone(), Some(90000)));
assert_eq!(frames.len(), 1);
assert_eq!(
frames[0].data, pcm,
"DVD must preserve every PCM byte — no second strip"
);
assert_eq!(frames[0].pts_ns, 1_000_000_000);
}
#[test]
fn dvd_lpcm_emits_short_payload_bd_would_drop() {
// A 4-byte raw-PCM DVD payload (1 sample pair at 16-bit stereo). The BD
// parser drops <= 4 bytes as "header only"; the DVD parser must emit it.
let mut bd = LpcmParser::new();
let mut dvd = LpcmParser::new_dvd();
let pcm = vec![0xAA, 0xBB, 0xCC, 0xDD];
assert!(
bd.parse(&make_pes(pcm.clone(), Some(0))).is_empty(),
"BD treats 4 bytes as header-only"
);
let frames = dvd.parse(&make_pes(pcm.clone(), Some(0)));
assert_eq!(frames.len(), 1);
assert_eq!(frames[0].data, pcm);
}
#[test]
fn always_keyframe() {
let mut parser = LpcmParser::new();
src/mux/codec/mod.rs
+11 -1
View File
@@ -102,7 +102,16 @@ impl CodecParser for PassthroughParser {
/// Create the appropriate parser for a codec, with optional codec private data.
///
/// For DvdSub, `codec_data` should be the pre-formatted VobSub .idx palette header.
pub fn parser_for_codec(codec: Codec, codec_data: Option<Vec<u8>>) -> Box<dyn CodecParser> {
///
/// `is_dvd_ps` selects the DVD program-stream variant where it matters: DVD
/// LPCM arrives with its private sub-header already stripped by the
/// `PsDemuxer`, so the LPCM parser must NOT strip the 4-byte BD LPCM header
/// again (that would drop one PCM sample pair per PES → progressive drift).
pub fn parser_for_codec(
codec: Codec,
codec_data: Option<Vec<u8>>,
is_dvd_ps: bool,
) -> Box<dyn CodecParser> {
match codec {
Codec::H264 => Box::new(h264::H264Parser::new()),
Codec::Hevc => Box::new(hevc::HevcParser::new()),
@@ -112,6 +121,7 @@ pub fn parser_for_codec(codec: Codec, codec_data: Option<Vec<u8>>) -> Box<dyn Co
Codec::DtsHdMa | Codec::DtsHdHr | Codec::Dts => Box::new(dts::DtsParser::new()),
Codec::TrueHd => Box::new(truehd::TrueHdParser::new()),
Codec::Pgs => Box::new(pgs::PgsParser::new()),
Codec::Lpcm if is_dvd_ps => Box::new(lpcm::LpcmParser::new_dvd()),
Codec::Lpcm => Box::new(lpcm::LpcmParser::new()),
Codec::DvdSub => Box::new(dvdsub::DvdSubParser::new(codec_data)),
_ => Box::new(PassthroughParser::new(true)),
src/mux/codec/mpeg2.rs
+74 -3
View File
@@ -99,6 +99,7 @@ impl CodecParser for Mpeg2Parser {
let pts_ns = pes.dts.or(pes.pts).map(pts_to_ns).unwrap_or(0);
let data = &pes.data;
let mut keyframe = false;
let mut has_picture = false;
// Scan for start codes in the elementary stream data.
let mut pos = 0;
@@ -176,6 +177,7 @@ impl CodecParser for Mpeg2Parser {
PICTURE_CODE => {
// Picture header: bytes after start code contain temporal_reference
// (10 bits) + picture_coding_type (3 bits).
has_picture = true;
if sc + 5 < data.len() {
let picture_coding_type = (data[sc + 5] >> 3) & 0x07;
if picture_coding_type == PICTURE_TYPE_I {
@@ -190,6 +192,23 @@ impl CodecParser for Mpeg2Parser {
}
}
// A PES that carried a sequence header but no picture start code is a
// parameter-set-only access unit: it has no coded picture to emit.
// Emitting it as a standalone keyframe would put bare sequence-header
// bytes into frame data with no picture. The sequence header is
// captured into codec_private above and is re-emitted in-band on the
// next real picture's PES, so dropping the empty access unit loses
// nothing. Mirrors how the H.264/HEVC parsers skip parameter-set-only
// access units.
//
// Conservative: only drop when this PES actually contained a sequence
// header and no picture. A PES with neither (e.g. a slice
// continuation) still passes through unchanged, preserving real
// keyframe detection.
if !has_picture && contains_seq_header(data) {
return Vec::new();
}
vec![Frame {
pts_ns,
keyframe,
@@ -241,6 +260,21 @@ fn parse_aspect_ratio(hdr: &[u8]) -> Option<(u8, u8)> {
Some(ASPECT_RATIOS[ar_code])
}
/// Returns true if `data` contains a sequence-header start code (00 00 01 B3).
fn contains_seq_header(data: &[u8]) -> bool {
let mut pos = 0;
while let Some(sc) = find_start_code(data, pos) {
if sc + 3 >= data.len() {
break;
}
if data[sc + 3] == SEQ_HEADER_CODE {
return true;
}
pos = sc + 4;
}
false
}
/// Find the position of the next start code (00 00 01) at or after `from`.
fn find_start_code(data: &[u8], from: usize) -> Option<usize> {
if data.len() < from + 3 {
@@ -429,15 +463,15 @@ mod tests {
assert!(has_ext, "codec_private should include sequence extension");
}
// --- I-frame with sequence header = keyframe ---
// --- sequence header + picture = keyframe ---
#[test]
fn sequence_header_implies_keyframe() {
fn sequence_header_with_picture_is_keyframe() {
let mut parser = Mpeg2Parser::new();
let mut data = Vec::new();
data.extend_from_slice(&make_seq_header(720, 480, 3, 4));
// Even without an explicit picture header, a sequence header implies I-frame.
data.extend_from_slice(&make_picture_header(PICTURE_TYPE_I));
data.extend_from_slice(&[0xFF; 16]);
let pes = make_pes(data, Some(0));
@@ -445,6 +479,43 @@ mod tests {
assert_eq!(frames.len(), 1);
assert!(frames[0].keyframe);
// codecPrivate is still captured.
assert!(parser.codec_private().is_some());
}
// --- parameter-set-only PES (seq header, no picture) emits no frame ---
#[test]
fn sequence_header_only_pes_emits_no_frame() {
let mut parser = Mpeg2Parser::new();
// A PES carrying only a sequence header (+ extension), no picture.
let mut data = Vec::new();
data.extend_from_slice(&make_seq_header(1920, 1080, 3, 4));
data.extend_from_slice(&[0x00, 0x00, 0x01, SEQ_EXT_CODE]);
data.extend_from_slice(&[0x14, 0x8A, 0x00, 0x01, 0x00, 0x00]);
let pes = make_pes(data, Some(0));
let frames = parser.parse(&pes);
// No coded picture → no frame emitted, but the sequence header is
// still captured for codecPrivate.
assert!(
frames.is_empty(),
"parameter-set-only PES should not emit a frame"
);
assert!(
parser.codec_private().is_some(),
"sequence header should still be captured into codec_private"
);
// A following picture-bearing PES emits the real keyframe.
let mut data2 = Vec::new();
data2.extend_from_slice(&make_picture_header(PICTURE_TYPE_I));
data2.extend_from_slice(&[0xFF; 16]);
let frames2 = parser.parse(&make_pes(data2, Some(3600)));
assert_eq!(frames2.len(), 1);
assert!(frames2[0].keyframe);
}
// --- PTS conversion ---
src/mux/codec/pgs.rs
+38
View File
@@ -18,6 +18,13 @@
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
const SEGMENT_PCS: u8 = 0x16;
// Upper bound on a pending display set's accumulated bytes. Real PGS
// display sets are small (a 1080p RLE bitmap plus palette is well under
// 1 MB); a stream that keeps appending non-PCS segments without ever
// emitting a PCS is malformed. Cap accumulation to bound memory and
// drop further appends until the next PCS resyncs the parser. Mirrors
// the MAX_*_BYTES / MAX_*_BUF caps in the DTS and AC-3 parsers.
const MAX_PGS_PENDING_BYTES: usize = 4 * 1024 * 1024;
// Offset within the PES payload at which number_of_composition_objects
// lives in a PCS: 3-byte segment header + 10 bytes of PCS fields
// (video_w/h, frame_rate, comp_num, comp_state, palette_update,
@@ -90,7 +97,12 @@ impl CodecParser for PgsParser {
// a pending display, append; otherwise emit as-is.
None => {
if let Some((_, ref mut buf)) = self.pending {
// Bound accumulation: a well-formed display set is small.
// Past the cap, drop further appends (malformed stream);
// the next PCS will take/replace `pending` and resync.
if buf.len() + pes.data.len() <= MAX_PGS_PENDING_BYTES {
buf.extend_from_slice(&pes.data);
}
} else {
out.push(Frame {
pts_ns,
@@ -180,6 +192,32 @@ mod tests {
assert!(data.windows(5).any(|w| w == [0x15, 0x00, 0x02, 0xAA, 0xBB]));
}
#[test]
fn pending_buffer_is_capped() {
let mut parser = PgsParser::new();
// Open a display set.
let _ = parser.parse(&make_pes(pcs_bytes(1), Some(90000)));
// Flood with non-PCS segments far exceeding the cap.
let chunk = vec![0x15u8; 256 * 1024]; // 256 KB ODS-like segment
let floods = (MAX_PGS_PENDING_BYTES / chunk.len()) + 32;
for _ in 0..floods {
let frames = parser.parse(&make_pes(chunk.clone(), Some(90000)));
assert!(frames.is_empty(), "non-PCS appends should not emit");
}
// The pending buffer must not have grown without bound.
let pending_len = parser.pending.as_ref().map(|(_, b)| b.len()).unwrap_or(0);
assert!(
pending_len <= MAX_PGS_PENDING_BYTES,
"pending buffer {pending_len} exceeded cap {MAX_PGS_PENDING_BYTES}"
);
// A following PCS still resyncs and emits the (capped) pending set.
let frames = parser.parse(&make_pes(pcs_bytes(0), Some(180000)));
assert_eq!(frames.len(), 1);
}
#[test]
fn codec_private_none() {
let parser = PgsParser::new();
src/mux/codec/truehd.rs
+12
View File
@@ -16,6 +16,12 @@ use super::{CodecParser, Frame, PesPacket, pts_to_ns};
/// Duration of one TrueHD access unit in nanoseconds (1/1200 second).
const AU_DURATION_NS: i64 = 833_333;
/// Hard cap on the reassembly buffer. A valid TrueHD/MAT access unit is
/// well under 32 KiB; if the buffer grows far past that without yielding a
/// frame the stream is malformed, so we drop it and resync rather than grow
/// without bound. Parity with the AC-3 / DTS / PGS caps.
const MAX_TRUEHD_BUF: usize = 256 * 1024;
pub struct TrueHdParser {
buf: Vec<u8>,
next_pts_ns: i64,
@@ -145,6 +151,12 @@ impl CodecParser for TrueHdParser {
self.next_pts_ns += AU_DURATION_NS;
}
// Bound memory on malformed input: a stream that never yields a
// complete frame must not grow the buffer without limit.
if self.buf.len() > MAX_TRUEHD_BUF {
self.buf.clear();
}
frames
}
src/mux/codec/vc1.rs
+53 -3
View File
@@ -44,6 +44,7 @@ impl CodecParser for Vc1Parser {
// Use DTS when available (monotonic for B-frame content), fall back to PTS
let ts_ns = pes.dts.or(pes.pts).map(pts_to_ns).unwrap_or(0);
let mut has_seq_header = false;
let mut has_entry_point = false;
let mut frame_start: Option<usize> = None;
// Scan for start codes (00 00 01 XX)
@@ -67,6 +68,7 @@ impl CodecParser for Vc1Parser {
SC_ENTRY_POINT => {
let end = find_next_sc(data, i + 4).unwrap_or(data.len());
self.entry_point = Some(data[i..end].to_vec());
has_entry_point = true;
}
SC_FRAME => {
// Frame data starts at this start code
@@ -85,11 +87,23 @@ impl CodecParser for Vc1Parser {
// Keyframe = this PES contains a sequence header (I-frame indicator in BD)
let keyframe = has_seq_header;
// Strip sequence header + entry point from frame data — those are in codecPrivate.
// Only include data from the frame start code onwards.
// Strip sequence header + entry point from frame data — those are in
// codecPrivate, not coded-picture data. Only include data from the
// frame start code onwards.
let frame_data = match frame_start {
Some(start) => &data[start..],
None => data, // no frame start code found, pass through entire PES
None => {
// No frame start code. If this PES carried only parameter sets
// (sequence header / entry point, captured above into
// codecPrivate), there is no coded picture to emit — drop it
// rather than passing parameter bytes through as a bogus
// keyframe. Mirrors how the H.264/HEVC parsers skip
// parameter-set-only access units.
if has_seq_header || has_entry_point {
return Vec::new();
}
data // genuine picture payload with no leading 0x0D — pass through
}
};
vec![Frame {
@@ -343,6 +357,42 @@ mod tests {
assert_eq!(&frames[0].data[0..4], &[0x00, 0x00, 0x01, SC_FRAME]);
}
// --- parameter-set-only PES (seq header + entry point, no frame SC) ---
#[test]
fn param_set_only_pes_emits_no_frame() {
let mut parser = Vc1Parser::new();
// Sequence header + entry point, but NO frame start code (0x0D).
let mut data = Vec::new();
data.extend_from_slice(&[0x00, 0x00, 0x01, SC_SEQUENCE_HEADER]);
data.extend_from_slice(&[0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF]);
data.extend_from_slice(&[0x00, 0x00, 0x01, SC_ENTRY_POINT]);
data.extend_from_slice(&[0x11, 0x22, 0x33, 0x44]);
let pes = make_pes(data, Some(90000));
let frames = parser.parse(&pes);
// No coded picture → no frame emitted (parameter bytes must not be
// passed through as a bogus keyframe).
assert!(
frames.is_empty(),
"parameter-set-only PES should not emit a frame"
);
// But codecPrivate is still captured.
assert!(parser.seq_header.is_some());
assert!(parser.entry_point.is_some());
assert!(parser.codec_private().is_some());
// A following frame-bearing PES still emits its picture.
let mut data2 = Vec::new();
data2.extend_from_slice(&[0x00, 0x00, 0x01, SC_FRAME]);
data2.extend_from_slice(&[0x55, 0x66, 0x77]);
let frames2 = parser.parse(&make_pes(data2, Some(180000)));
assert_eq!(frames2.len(), 1);
assert_eq!(&frames2[0].data[0..4], &[0x00, 0x00, 0x01, SC_FRAME]);
}
// --- empty PES ---
#[test]
src/mux/disc.rs
+2 -1
View File
@@ -195,7 +195,8 @@ impl DiscStream {
};
pids.push(pid);
pid_to_track.push((pid, idx));
parsers.push((pid, super::codec::parser_for_codec(codec, None)));
let is_dvd_ps = matches!(content_format, crate::disc::ContentFormat::MpegPs);
parsers.push((pid, super::codec::parser_for_codec(codec, None, is_dvd_ps)));
}
let mut ts_demuxer = None;
src/mux/ebml.rs
+38 -7
View File
@@ -264,6 +264,13 @@ pub fn read_element_header(r: &mut impl Read) -> io::Result<(u32, u64, usize)> {
/// Read an unsigned integer value of `len` bytes.
pub fn read_uint_val(r: &mut impl Read, len: usize) -> io::Result<u64> {
// An EBML unsigned integer is at most 8 bytes. A malformed element
// claiming `len > 8` would index past this stack buffer and panic
// (DoS on untrusted input) — reject it at the source so every caller
// is safe, not just the ones that pre-check.
if len > 8 {
return Err(crate::error::Error::MkvInvalid.into());
}
let mut buf = [0u8; 8];
r.read_exact(&mut buf[..len])?;
let mut val = 0u64;
@@ -273,23 +280,33 @@ pub fn read_uint_val(r: &mut impl Read, len: usize) -> io::Result<u64> {
Ok(val)
}
/// Read a float value (4 or 8 bytes).
/// Read a float value. EBML floats are exactly 0, 4, or 8 bytes.
///
/// The previous `else` branch read a fixed 8 bytes for ANY non-4 length,
/// so a malformed element with `len > 8` left `len - 8` unconsumed bytes
/// (mis-read as the next EBML header → desync of the rest of the parent
/// element) and `len < 4` over-read. Consume exactly `len` bytes and
/// reject anything that isn't a valid float width.
pub fn read_float_val(r: &mut impl Read, len: usize) -> io::Result<f64> {
if len == 4 {
match len {
0 => Ok(0.0),
4 => {
let mut buf = [0u8; 4];
r.read_exact(&mut buf)?;
Ok(f32::from_be_bytes(buf) as f64)
} else {
}
8 => {
let mut buf = [0u8; 8];
r.read_exact(&mut buf)?;
Ok(f64::from_be_bytes(buf))
}
_ => Err(crate::error::Error::MkvInvalid.into()),
}
}
/// Read a UTF-8 string value of `len` bytes.
pub fn read_string_val(r: &mut impl Read, len: usize) -> io::Result<String> {
let mut buf = vec![0u8; len];
r.read_exact(&mut buf)?;
let mut buf = read_exact_bounded(r, len)?;
// Strip trailing nulls
while buf.last() == Some(&0) {
buf.pop();
@@ -299,8 +316,22 @@ pub fn read_string_val(r: &mut impl Read, len: usize) -> io::Result<String> {
/// Read binary data of `len` bytes.
pub fn read_binary_val(r: &mut impl Read, len: usize) -> io::Result<Vec<u8>> {
let mut buf = vec![0u8; len];
r.read_exact(&mut buf)?;
read_exact_bounded(r, len)
}
/// Read exactly `len` bytes WITHOUT trusting `len` to size the allocation.
///
/// `vec![0u8; len]` on an attacker-controlled EBML size would allocate
/// gigabytes before the read fails. Instead we cap the reader to `len`
/// and grow the buffer as bytes actually arrive: a malformed element that
/// claims a huge length but supplies few bytes allocates only what it
/// delivers, then errors on the short read.
fn read_exact_bounded(r: &mut impl Read, len: usize) -> io::Result<Vec<u8>> {
let mut buf = Vec::new();
let got = r.take(len as u64).read_to_end(&mut buf)?;
if got != len {
return Err(io::ErrorKind::UnexpectedEof.into());
}
Ok(buf)
}
src/mux/mkv.rs
+145 -10
View File
@@ -112,19 +112,22 @@ impl MkvTrack {
}
pub fn audio(a: &AudioStream) -> Self {
// Codec ID strings must distinguish the DTS family — strict
// players (Plex transcoder, some hardware decoders, some AV
// receivers) reject lossless DTS-HD MA payload when the
// track advertises plain `A_DTS` because it implies the
// bitstream is the 1.5 Mbps "core" only. Fix is just to
// emit the right ID per BD-STN codec field.
// The Matroska codec-ID registry defines `A_DTS` for the entire
// DTS family — the spec text for `A_DTS` explicitly states it
// "Supports DTS, DTS-ES, DTS-96/26, DTS-HD High Resolution Audio
// and DTS-HD Master Audio." Players distinguish core vs HD-HRA vs
// HD-MA by parsing the DTS bitstream extension substreams, not by
// the container codec ID. The previously-emitted `A_DTS/MA` and
// `A_DTS/HR` suffixes are NOT registered codec IDs; strict parsers
// (libmatroska) and some hardware renderers fail to recognise the
// track at all. Emit plain `A_DTS` for every DTS variant — the
// lossless MA / HRA payload bytes are unchanged, only the
// container codec-ID string differs.
let codec_id = match a.codec {
Codec::Ac3 => "A_AC3",
Codec::Ac3Plus => "A_EAC3",
Codec::TrueHd => "A_TRUEHD",
Codec::DtsHdMa => "A_DTS/MA",
Codec::DtsHdHr => "A_DTS/HR",
Codec::Dts => "A_DTS",
Codec::DtsHdMa | Codec::DtsHdHr | Codec::Dts => "A_DTS",
Codec::Lpcm => "A_PCM/INT/BIG",
_ => "A_AC3",
};
@@ -220,6 +223,12 @@ pub struct MkvMuxer<W: Write + Seek> {
/// New cluster every 5 seconds.
const CLUSTER_DURATION_MS: i64 = 5000;
/// Maximum block-relative timestamp expressible in the signed 16-bit
/// SimpleBlock/Block field (`i16::MAX` ms). A frame further than this from
/// the open cluster's timestamp forces a new cluster (see `write_frame`) so
/// the `as i16` cast can never wrap.
const MAX_BLOCK_REL_MS: i64 = i16::MAX as i64;
impl<W: Write + Seek> MkvMuxer<W> {
/// Create a new MKV muxer: writes EBML header, Segment start, Info, Tracks, Chapters.
pub fn new(
@@ -444,7 +453,7 @@ impl<W: Write + Seek> MkvMuxer<W> {
let base = *self.base_pts_ms.get_or_insert(raw_ms);
let pts_ms = raw_ms - base;
// Cluster boundaries must coincide with a video keyframe so every
// Cluster boundaries normally coincide with a video keyframe so every
// Cues entry resolves to a seekable IDR at the cluster start.
let is_video_key = keyframe && track_idx == 0;
let needs_new_cluster = !self.cluster_open
@@ -460,6 +469,17 @@ impl<W: Write + Seek> MkvMuxer<W> {
track: track_idx + 1,
cluster_pos: self.cluster_pos - self.segment_start,
});
} else if (pts_ms - self.cluster_ts_ms) > MAX_BLOCK_REL_MS {
// The block-relative timestamp is a signed 16-bit value, so a
// frame more than i16::MAX ms (~32.767 s) past the current
// cluster's timestamp would silently wrap on the `as i16` cast,
// corrupting A/V sync. The keyframe-driven boundary above only
// fires on a video keyframe — a long audio-only stretch, or a
// very long GOP with no intervening keyframe, can drift past the
// i16 range. Force a fresh cluster here even without a keyframe
// to keep the cast in range. This cluster is not keyframe-aligned
// so it gets no Cues entry (Cues stay IDR-only for seekability).
self.start_cluster(pts_ms)?;
}
let relative_ts = (pts_ms - self.cluster_ts_ms) as i16;
@@ -679,6 +699,42 @@ mod tests {
}
}
fn audio_stream(codec: Codec) -> AudioStream {
use crate::disc::{AudioChannels, LabelPurpose, SampleRate};
AudioStream {
pid: 0x1100,
codec,
channels: AudioChannels::Surround51,
language: "eng".into(),
sample_rate: SampleRate::S48,
secondary: false,
purpose: LabelPurpose::Normal,
label: String::new(),
}
}
#[test]
fn dts_variants_map_to_registered_a_dts_codec_id() {
// The Matroska codec-ID registry defines `A_DTS` for the whole DTS
// family (core, DTS-HD HRA, DTS-HD MA). The `/MA` and `/HR` suffixes
// are not registered and break strict parsers, so every DTS variant
// must emit plain `A_DTS`.
for codec in [Codec::Dts, Codec::DtsHdMa, Codec::DtsHdHr] {
let track = MkvTrack::audio(&audio_stream(codec));
assert_eq!(
track.codec_id, "A_DTS",
"{codec:?} must map to registered codec ID A_DTS, got {}",
track.codec_id
);
}
// Sanity: the non-DTS variants keep their distinct IDs.
assert_eq!(MkvTrack::audio(&audio_stream(Codec::Ac3)).codec_id, "A_AC3");
assert_eq!(
MkvTrack::audio(&audio_stream(Codec::TrueHd)).codec_id,
"A_TRUEHD"
);
}
#[test]
fn dolby_vision_config_profile7() {
// dvcC for disc Profile 7 dual-layer: version 1.0, profile 7, all of
@@ -1377,6 +1433,85 @@ mod tests {
);
}
/// Collect every (cluster_ts_ms, block_relative_ts_i16, absolute_ms) for
/// all SimpleBlocks across all clusters, so a test can assert that the
/// reconstructed absolute timestamp (cluster_ts + relative_ts) is correct
/// and that no relative_ts ever wrapped the i16 range.
fn all_block_timestamps(data: &[u8]) -> Vec<(i64, i16, i64)> {
let mut out = Vec::new();
for (body_start, body_size, cluster_ts) in find_clusters(data) {
let body = &data[body_start..body_start + body_size as usize];
let mut cursor = Cursor::new(body);
// Skip CLUSTER_TIMESTAMP.
let (tid, tsize, _) = ebml::read_element_header(&mut cursor).unwrap();
assert_eq!(tid, ebml::CLUSTER_TIMESTAMP);
cursor.seek(io::SeekFrom::Current(tsize as i64)).unwrap();
while (cursor.position() as usize) < body.len() {
let (id, sz, _) = ebml::read_element_header(&mut cursor).unwrap();
if id == ebml::SIMPLE_BLOCK {
let bstart = cursor.position() as usize;
let b0 = body[bstart];
let vint_len = if b0 & 0x80 != 0 { 1 } else { 2 };
let ts_pos = bstart + vint_len;
let rel = i16::from_be_bytes([body[ts_pos], body[ts_pos + 1]]);
out.push((cluster_ts as i64, rel, cluster_ts as i64 + rel as i64));
}
cursor.seek(io::SeekFrom::Current(sz as i64)).unwrap();
}
}
out
}
#[test]
fn long_audio_gap_forces_cluster_no_i16_overflow() {
// Regression for the `(pts_ms - cluster_ts_ms) as i16` truncation:
// a single video keyframe at t=0 opens one cluster, then a long
// audio-only stretch (no further video keyframe) drifts well past
// i16::MAX ms (~32.767 s). Without the overflow guard the audio
// blocks past 32.767 s would write a wrapped (negative) relative
// timestamp into the SimpleBlock. With the guard a fresh cluster is
// forced so every relative_ts stays in range and reconstructs to the
// true absolute timestamp.
let tracks = [make_video_track(), make_audio_track()];
let mut frames: Vec<(usize, i64, bool, Vec<u8>)> = Vec::new();
// One video keyframe at t=0 (opens the first cluster).
frames.push((0, 0, true, vec![0xAB; 16]));
// Audio frames every 100 ms out to 60 s — past the 32.767 s i16 limit
// and past two i16 spans, with NO further video keyframe.
let mut t_ms = 0i64;
while t_ms <= 60_000 {
frames.push((1, t_ms * 1_000_000, true, vec![0xCD; 16]));
t_ms += 100;
}
let (data, _) = mux_to_bytes(&tracks, &[], &frames);
let blocks = all_block_timestamps(&data);
assert!(!blocks.is_empty());
// Every block's relative timestamp must be within i16 range (it is by
// type), AND must reconstruct to a non-negative, monotonic-ish
// absolute timestamp matching the source — i.e. no silent wrap.
for (cluster_ts, rel, abs) in &blocks {
assert!(
*rel as i64 >= 0 && (*rel as i64) <= MAX_BLOCK_REL_MS,
"block relative_ts {rel} out of [0, i16::MAX] range \
(cluster_ts={cluster_ts}, abs={abs}) — i16 overflow"
);
}
// The latest audio frame is at 60_000 ms; its reconstructed absolute
// timestamp must equal that, proving no truncation occurred.
let max_abs = blocks.iter().map(|(_, _, abs)| *abs).max().unwrap();
assert_eq!(max_abs, 60_000, "last block must reconstruct to 60_000 ms");
// The overflow guard must have opened more than one cluster (the
// single keyframe alone would otherwise yield exactly one).
let clusters = find_clusters(&data);
assert!(
clusters.len() >= 2,
"expected the i16 guard to force extra clusters, got {}",
clusters.len()
);
}
#[test]
fn pre_first_keyframe_frames_dropped() {
let tracks = [make_video_track()];
src/mux/mkvstream.rs
+203 -16
View File
@@ -14,6 +14,49 @@ fn skip_bytes(r: &mut impl Read, n: u64) -> io::Result<()> {
Ok(())
}
// ── Sanity caps for untrusted EBML element sizes ──────────────
//
// Sizes come straight from the EBML stream (file or network) and are
// otherwise cast to `usize` and used to allocate/read. An adversarial
// or corrupt container can claim a multi-GB element and trigger an OOM
// allocation, or claim an integer element wider than 8 bytes and panic
// the fixed 8-byte reader. Every untrusted size is validated against
// one of these caps before allocation.
/// Largest accepted SIMPLE_BLOCK payload. A block is a small vint track
/// header + 2-byte rel-ts + 1-byte flags + one frame of elementary data.
/// UHD HEVC keyframes run a few MB; 64 MiB is generously above any real
/// single-frame block while still bounding a hostile allocation.
const MAX_BLOCK_SIZE: u64 = 64 * 1024 * 1024;
/// Largest accepted CODEC_PRIVATE payload. hvcC/avcC/setup blobs are a
/// few KB in practice; 16 MiB is far above any legitimate value.
const MAX_CODEC_PRIVATE: u64 = 16 * 1024 * 1024;
/// Largest accepted string element (TITLE/CODEC_ID/LANGUAGE/TRACK_NAME).
const MAX_STRING_LEN: u64 = 64 * 1024;
/// EBML unsigned-int elements are at most 8 bytes wide.
const MAX_UINT_LEN: u64 = 8;
/// Reject an untrusted element size that exceeds `cap` before it is used
/// to allocate or read. Returns the size as `usize` when within bounds.
fn checked_size(size: u64, cap: u64) -> io::Result<usize> {
if size > cap {
return Err(crate::error::Error::MkvInvalid.into());
}
Ok(size as usize)
}
/// Read a bounded unsigned int. Guards against `size > 8` (which would
/// otherwise index out of the fixed 8-byte buffer in `read_uint_val`)
/// before delegating.
fn read_uint_bounded(r: &mut impl Read, size: u64) -> io::Result<u64> {
ebml::read_uint_val(r, checked_size(size, MAX_UINT_LEN)?)
}
/// Read a bounded UTF-8 string element.
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};
@@ -108,11 +151,12 @@ impl crate::pes::Stream for MkvStream {
match id {
ebml::CLUSTER => continue,
ebml::CLUSTER_TIMESTAMP => {
rs.cluster_ts_ms = ebml::read_uint_val(&mut rs.reader, size as usize)? as i64;
rs.cluster_ts_ms = read_uint_bounded(&mut rs.reader, size)? as i64;
continue;
}
ebml::SIMPLE_BLOCK => {
let block = ebml::read_binary_val(&mut rs.reader, size as usize)?;
let block =
ebml::read_binary_val(&mut rs.reader, checked_size(size, MAX_BLOCK_SIZE)?)?;
if block.len() < 4 {
continue;
}
@@ -198,7 +242,10 @@ impl crate::pes::Stream for MkvStream {
/// Returns (DiscTitle, codec_privates: Vec<(track_number, codec_private_bytes)>)
fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult {
let mut title = String::new();
let mut duration_ms = 0.0f64;
// EBML `DURATION` is a float expressed in TimestampScale ticks, not
// milliseconds (Matroska spec). Named accordingly; converted to
// seconds below as ticks * ts_scale_ns / 1e9.
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();
@@ -235,9 +282,9 @@ fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult {
let (cid, cs, hlen) = ebml::read_element_header(r)?;
remaining = remaining.saturating_sub(hlen as u64 + cs);
match cid {
ebml::TIMESTAMP_SCALE => ts_scale = ebml::read_uint_val(r, cs as usize)?,
ebml::DURATION => duration_ms = ebml::read_float_val(r, cs as usize)?,
ebml::TITLE => title = ebml::read_string_val(r, cs as usize)?,
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)?;
}
@@ -274,7 +321,7 @@ fn parse_mkv_header(r: &mut impl Read) -> MkvHeaderResult {
let disc_title = DiscTitle {
playlist: title,
duration_secs: duration_ms * (ts_scale as f64) / 1_000_000_000.0,
duration_secs: duration_ticks * (ts_scale as f64) / 1_000_000_000.0,
streams,
..DiscTitle::empty()
};
@@ -296,20 +343,25 @@ fn parse_track(
let (cid, cs, hlen) = ebml::read_element_header(r)?;
remaining = remaining.saturating_sub(hlen as u64 + cs);
match cid {
ebml::TRACK_NUMBER => tnum = ebml::read_uint_val(r, cs as usize)? as u16,
ebml::TRACK_TYPE => ttype = ebml::read_uint_val(r, cs as usize)?,
ebml::CODEC_ID => codec_id = ebml::read_string_val(r, cs as usize)?,
ebml::CODEC_PRIVATE => codec_priv = Some(ebml::read_binary_val(r, cs as usize)?),
ebml::LANGUAGE => lang = ebml::read_string_val(r, cs as usize)?,
ebml::TRACK_NAME => name = ebml::read_string_val(r, cs as usize)?,
ebml::FLAG_FORCED => forced = ebml::read_uint_val(r, cs as usize)? != 0,
ebml::TRACK_NUMBER => tnum = read_uint_bounded(r, cs)? 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)?;
vrem = vrem.saturating_sub(vhlen as u64 + vs);
if vid == ebml::PIXEL_HEIGHT {
ph = ebml::read_uint_val(r, vs as usize)? as u32;
ph = read_uint_bounded(r, vs)? as u32;
} else {
skip_bytes(r, vs)?;
}
@@ -322,7 +374,7 @@ fn parse_track(
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 = ebml::read_uint_val(r, as_ as usize)? as u8,
ebml::CHANNELS => ch = read_uint_bounded(r, as_)? as u8,
_ => {
skip_bytes(r, as_)?;
}
@@ -428,3 +480,138 @@ fn block_vint(d: &[u8]) -> (u64, usize) {
}
(0, 1) // Unsupported 5+ byte VINT — treat as track 0
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pes::Stream as _;
use std::io::Cursor;
// `From<Error> for io::Error` encodes the numeric code into the
// Display string as "E{code}: ...". Check the prefix.
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 checked_size_rejects_over_cap() {
// Within cap → Ok with usize value.
assert_eq!(checked_size(100, 256).unwrap(), 100);
assert_eq!(checked_size(256, 256).unwrap(), 256);
// Over cap → MkvInvalid, never a giant allocation.
let e = checked_size(257, 256).unwrap_err();
assert!(is_mkv_invalid(&e));
// A hostile multi-GB block size is rejected as MkvInvalid.
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() {
// size > 8 would index out of the fixed 8-byte buffer in
// read_uint_val (panic / OOB). The guard turns it into a clean
// MkvInvalid error instead.
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() {
// 8 bytes is the max legal EBML uint width and must still work.
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() {
// Claimed string length far above the cap must not allocate.
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));
}
/// Build a minimal MKV (EBML header + Segment + Info + Tracks) so the
/// reader is positioned in the cluster body, then append the given
/// cluster bytes. Returns the full byte stream.
fn minimal_mkv_with_cluster(cluster_body: &[u8]) -> Vec<u8> {
let mut out = Vec::new();
// EBML header (empty body).
ebml::write_id(&mut out, ebml::EBML).unwrap();
ebml::write_size(&mut out, 0).unwrap();
// Segment (unknown size so the reader streams children).
ebml::write_id(&mut out, ebml::SEGMENT).unwrap();
ebml::write_unknown_size(&mut out).unwrap();
// Empty Info.
ebml::write_id(&mut out, ebml::INFO).unwrap();
ebml::write_size(&mut out, 0).unwrap();
// Empty Tracks.
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() {
// Cluster containing a SIMPLE_BLOCK that claims a 2 GiB payload.
// The reader must reject it (MkvInvalid) rather than attempt a
// multi-GB allocation. Header parse stops at CLUSTER, so the
// SIMPLE_BLOCK is hit on the first read().
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();
// No payload follows — but we must fail on the size check, before
// any read of the body.
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() {
// A small, well-formed SIMPLE_BLOCK must still parse into a frame.
// We need at least one stream so the track index is in range, so
// give Tracks one video TRACK_ENTRY (track number 1).
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();
// Tracks → one TRACK_ENTRY (track number 1, type 1 = video).
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);
// Cluster with a SIMPLE_BLOCK: track vint=0x81 (track 1),
// rel_ts=0x0000, flags=0x80 (keyframe), then 4 bytes of data.
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]);
}
}
src/mux/resolve.rs
+2 -1
View File
@@ -384,7 +384,8 @@ fn build_demux_state(title: &DiscTitle, format: ContentFormat) -> DemuxState {
};
pids.push(pid);
pid_to_track.push((pid, idx));
parsers.push((pid, super::codec::parser_for_codec(codec, None)));
let is_dvd_ps = matches!(format, ContentFormat::MpegPs);
parsers.push((pid, super::codec::parser_for_codec(codec, None, is_dvd_ps)));
}
let (ts, ps) = match format {
ContentFormat::MpegPs => (None, Some(super::ps::PsDemuxer::new())),
src/mux/ts.rs
+59 -7
View File
@@ -35,6 +35,15 @@ struct PesAssembler {
pts: Option<i64>,
dts: Option<i64>,
active: bool,
/// PES-header bytes still to be skipped on the next continuation
/// packet(s). A PES header (9 + PES_header_data_length, up to 264
/// bytes) can exceed a single 184-byte TS payload, spilling into the
/// following continuation packet. Those spillover bytes are NOT
/// elementary-stream data and must be skipped, or the PES start code
/// (`00 00 01 …`) and timestamp bytes get injected into the ES — for
/// HEVC/H264 that reads as a spurious start code / corrupt slice
/// payload. Tracks how many header bytes remain across packets.
header_remaining: usize,
}
/// Initial capacity for a fresh PES buffer. Sized to cover the
@@ -54,6 +63,7 @@ impl PesAssembler {
pts: None,
dts: None,
active: false,
header_remaining: 0,
}
}
@@ -235,12 +245,35 @@ impl TsDemuxer {
let payload = &ts[payload_start..];
if pusi {
let (pts, dts, pes_data_start) = parse_pes_header(payload);
// `header_len` is the FULL (uncapped) PES-header length:
// 0 = malformed (payload is not a PES start), else 6/9+N.
let (pts, dts, header_len) = parse_pes_header(payload);
if let Some(prev) = asm.start(pts, dts) {
completed.push(prev);
}
if pes_data_start < payload.len() {
asm.push(&payload[pes_data_start..]);
if header_len == 0 {
// PUSI packet whose payload is not a valid PES start. Do
// NOT push it — those bytes are not elementary-stream data
// and would inject a spurious start code / garbage.
asm.header_remaining = 0;
} else if header_len <= payload.len() {
// Header fits in this packet (the common case).
asm.header_remaining = 0;
if header_len < payload.len() {
asm.push(&payload[header_len..]);
}
} else {
// Header spills past this packet — skip the remainder on
// the following continuation packet(s).
asm.header_remaining = header_len - payload.len();
}
} else if asm.header_remaining > 0 {
// Continuation packet still inside a PES header that spanned
// the boundary — consume header bytes before any ES data.
let skip = asm.header_remaining.min(payload.len());
asm.header_remaining -= skip;
if skip < payload.len() {
asm.push(&payload[skip..]);
}
} else {
asm.push(payload);
@@ -260,7 +293,14 @@ impl TsDemuxer {
}
/// Parse a PES packet header, extracting PTS and DTS.
/// Returns (pts, dts, offset_to_elementary_stream_data).
///
/// Returns `(pts, dts, header_len)` where `header_len` is the FULL,
/// UNCAPPED PES-header length in bytes (`9 + PES_header_data_length`, or
/// 6 for stream IDs without the standard extension). `0` signals the
/// payload is not a valid PES start (malformed / too short). The caller
/// must treat `header_len` as bytes-to-skip and carry any remainder past
/// this packet's payload into the next continuation packet — the header
/// can exceed one TS payload, and the spillover is header, not ES data.
fn parse_pes_header(data: &[u8]) -> (Option<i64>, Option<i64>, usize) {
// PES packet: 00 00 01 [stream_id] [length:2] [flags...]
if data.len() < 9 || data[0] != 0x00 || data[1] != 0x00 || data[2] != 0x01 {
@@ -288,7 +328,10 @@ fn parse_pes_header(data: &[u8]) -> (Option<i64>, Option<i64>, usize) {
let pts_dts_flags = (data[7] >> 6) & 0x03;
let header_data_len = data[8] as usize;
let data_start = (9 + header_data_len).min(data.len());
// Full, uncapped header length. PTS/DTS (if present) live in the
// first ~19 bytes, always within this packet's payload, so they parse
// here; only the *skip* length may extend into the next packet.
let header_len = 9 + header_data_len;
let mut pts = None;
let mut dts = None;
@@ -300,7 +343,7 @@ fn parse_pes_header(data: &[u8]) -> (Option<i64>, Option<i64>, usize) {
dts = parse_timestamp(&data[14..19]);
}
(pts, dts, data_start)
(pts, dts, header_len)
}
/// Parse a 5-byte PTS/DTS timestamp (33 bits in 90kHz).
@@ -404,10 +447,19 @@ pub fn scan_streams(data: &[u8]) -> Option<Vec<crate::disc::Stream>> {
let section_len =
(((data[pmt_start + 1] & 0x0F) as usize) << 8) | data[pmt_start + 2] as usize;
// section_length counts the bytes after this field, including the
// trailing 4-byte CRC; `< 4` would underflow `end` below. Guard it
// exactly like the PAT parser above.
if section_len < 4 {
offset += BD_TS_PACKET_SIZE;
continue;
}
let prog_info_len =
(((data[pmt_start + 10] & 0x0F) as usize) << 8) | data[pmt_start + 11] as usize;
let mut pos = pmt_start + 12 + prog_info_len;
let end = pmt_start + 3 + section_len - 4;
// Clamp the section end to the buffer; a malformed section_len or
// prog_info_len must never drive reads past `data`.
let end = (pmt_start + 3 + section_len - 4).min(data.len());
while pos + 5 <= data.len() && pos < end {
let stream_type = data[pos];
src/sector/prefetched.rs
+8 -7
View File
@@ -157,14 +157,15 @@ impl PrefetchedSectorSource {
Ok(b) => b,
Err(_) => return, // consumer dropped both channels
};
if buf.len() < bytes {
buf.resize(bytes, 0);
} else {
// Re-expose the full extent; previous truncate
// shrank the visible len without freeing pages.
// SAFETY: capacity is at least `bytes` after
// construction with `vec![0u8; batch_bytes]`.
if bytes <= buf.capacity() {
// Re-expose `bytes` without zero-filling pages that
// `read_sectors` is about to overwrite. The capacity
// guard makes the `set_len` provably sound even if a
// recycled buffer ever comes back smaller than the
// `vec![0u8; batch_bytes]` it was born with.
unsafe { buf.set_len(bytes) };
} else {
buf.resize(bytes, 0);
}
let lba = extent.start_lba + offset;
match reader.read_sectors(lba, sectors, &mut buf[..bytes], false) {