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

0.31.0: hardening and correctness pass across mux, codec, AACS/CSS, UDF/MPLS/CLPI, recovery, drive/SCSI, labels, and I/O

Browse Source 
Library-wide review-and-fix pass: tightened AACS keydb/handshake/variant
handling and trailing-partial-unit policy, corrected MPLS mark offset and
added UDF allocation bounds, hardened the mux/codec framing and M2TS paths,
guarded SCSI READ CAPACITY short transfers and unified error mapping, added
overflow guards on untrusted disc input, and made prefetch shutdown
deterministic. Release profile now builds with thin LTO + single codegen unit.
This commit is contained in:
Matthew Jackson
2026-06-07 17:37:38 -07:00
parent 5b6ea8f5c4
commit 061f68594a
128 changed files with 11838 additions and 3831 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/mux/tsmux.rs
+257 -87
View File
@@ -4,17 +4,47 @@
//! packets. Each frame is wrapped in a PES header, split into TS packets,
//! and prepended with the 4-byte TP_extra_header.
use super::hevc::{hvcc_to_annex_b, length_prefixed_to_annex_b};
use std::io::{self, Write};
const SYNC_BYTE: u8 = 0x47;
const TS_PAYLOAD: usize = 184;
/// PID range treated as video (HEVC, triggers Annex-B conversion + RAI
/// on keyframes). Both `write_frame` and `build_pes_header` consult this
/// so a PID's stream_id and its NAL handling can never disagree.
const VIDEO_PID_RANGE: std::ops::RangeInclusive<u16> = 0x1011..=0x101F;
/// Largest PES payload that fits a bounded `PES_packet_length` (u16) on a
/// `0xBD` (private_stream_1) stream after the 8 PES-header bytes. Frames
/// larger than this are split into multiple PES so the length field stays
/// spec-conformant (the unbounded `0` length is only legal for video).
const MAX_BD_PES_PAYLOAD: usize = u16::MAX as usize - 8;
fn is_video_pid(pid: u16) -> bool {
VIDEO_PID_RANGE.contains(&pid)
}
/// BD-TS muxer: PES frames in, 192-byte BD-TS packets out.
///
/// Constructed over an output writer and a slice of per-track PIDs. The
/// `track` index passed to [`TsMuxer::write_frame`] and
/// [`TsMuxer::set_codec_private`] is the position in that PID slice; all
/// per-track state vectors are sized to `pids.len()`. PIDs in
/// `0x1011..=0x101F` are treated as video (length-prefixed NALUs in,
/// Annex B out, with parameter-set prepend and RAI on keyframes); every
/// other PID is carried as `private_stream_1` (`0xBD`) audio/subtitle.
/// All tracks share one PTS origin seeded from the first video frame, so
/// audio/video PTS offsets are preserved.
pub struct TsMuxer<W: Write> {
writer: W,
pids: Vec<u16>,
continuity: Vec<u8>, // per-PID continuity counter (0-15)
codec_privates: Vec<Option<Vec<u8>>>, // per-track codec_private (for video parameter sets)
params_written: Vec<bool>, // per-track: have we written parameter sets?
/// Global PTS origin (nanoseconds), seeded by the FIRST video frame so
/// the audio/video offset is preserved. Frames that arrive before it
/// is set saturate to 0.
base_pts_ns: Option<i64>,
}
@@ -33,15 +63,29 @@ impl<W: Write> TsMuxer<W> {
/// Set codec_private data for a track. Used to prepend VPS/SPS/PPS
/// as Annex B NALs before the first keyframe in the transport stream.
pub fn set_codec_private(&mut self, track: usize, data: Vec<u8>) {
if track < self.codec_privates.len() {
self.codec_privates[track] = Some(data);
///
/// `track` is the index into the PID slice passed to [`TsMuxer::new`].
/// Returns [`Error::MuxTrackRange`](crate::error::Error::MuxTrackRange)
/// for an out-of-range index.
pub fn set_codec_private(&mut self, track: usize, data: Vec<u8>) -> io::Result<()> {
if track >= self.codec_privates.len() {
return Err(crate::error::Error::MuxTrackRange {
track,
tracks: self.codec_privates.len(),
}
.into());
}
self.codec_privates[track] = Some(data);
Ok(())
}
/// Write a PES frame as BD-TS packets.
/// Video frame data is expected as length-prefixed NALUs (MKV/PES format)
/// and is converted to Annex B for transport stream.
///
/// `track` is the index into the PID slice passed to [`TsMuxer::new`].
/// Returns [`Error::MuxTrackRange`](crate::error::Error::MuxTrackRange)
/// for an out-of-range index.
pub fn write_frame(
&mut self,
track: usize,
@@ -50,10 +94,14 @@ impl<W: Write> TsMuxer<W> {
data: &[u8],
) -> io::Result<()> {
if track >= self.pids.len() {
return Ok(()); // unknown track, skip
return Err(crate::error::Error::MuxTrackRange {
track,
tracks: self.pids.len(),
}
.into());
}
let pid = self.pids[track];
let is_video = (0x1011..=0x101F).contains(&pid);
let is_video = is_video_pid(pid);
// Drop non-key video before any keyframe — decoder has no IDR or
// parameter sets to anchor on.
@@ -61,12 +109,26 @@ impl<W: Write> TsMuxer<W> {
return Ok(());
}
let base = *self.base_pts_ns.get_or_insert(pts_ns);
let pts_ns = pts_ns - base;
// Seed the global PTS origin from the FIRST video frame only, so the
// audio/video offset is preserved. A leading audio frame must not
// pull the base up and collapse the first video IDR to t=0.
if is_video {
self.base_pts_ns.get_or_insert(pts_ns);
}
let base = self.base_pts_ns.unwrap_or(pts_ns);
let pts_ns = pts_ns.saturating_sub(base);
// For video: convert length-prefixed NALUs to Annex B (start codes).
// Prepend codec_private parameter sets on the FIRST keyframe only.
let es_data = if is_video && !data.is_empty() {
//
// Arm `params_written` on the first video keyframe regardless of
// whether it carries data: an empty-data keyframe still anchors
// the stream, and leaving the flag unset would make every later
// non-key frame fail the drop guard above and silently vanish.
// For non-video the ES bytes pass through unchanged, so borrow
// `data` directly rather than copying it; only video needs an
// owned Annex-B conversion buffer.
let es_data: std::borrow::Cow<'_, [u8]> = if is_video {
let mut annex_b = Vec::new();
if keyframe && !self.params_written[track] {
if let Some(ref cp) = self.codec_privates[track] {
@@ -77,25 +139,58 @@ impl<W: Write> TsMuxer<W> {
self.params_written[track] = true;
}
annex_b.extend_from_slice(&length_prefixed_to_annex_b(data));
annex_b
std::borrow::Cow::Owned(annex_b)
} else {
data.to_vec()
std::borrow::Cow::Borrowed(data)
};
// Build PES packet: header + data
let pts_90k = if pts_ns >= 0 {
(pts_ns as u64).saturating_mul(9) / 100_000
} else {
0
};
let pes_header = build_pes_header(pid, pts_90k, es_data.len());
let pes_packet = [&pes_header[..], &es_data[..]].concat();
// Split into TS packets
// Video PES may be unbounded (length 0); a 0xBD private_stream_1
// PES must carry a bounded length, so split oversized audio/sub
// access units into multiple PES packets. Each emitted PES carries
// the same PTS and starts on its own PUSI packet (only the keyframe
// RAI rides the first packet of the first PES).
if is_video || es_data.len() <= MAX_BD_PES_PAYLOAD {
self.write_pes_chain(track, pid, pts_90k, is_video, keyframe, &es_data)?;
} else {
let mut first_pes = true;
for chunk in es_data.chunks(MAX_BD_PES_PAYLOAD) {
self.write_pes_chain(track, pid, pts_90k, is_video, keyframe && first_pes, chunk)?;
first_pes = false;
}
}
Ok(())
}
/// Wrap `es_data` in a PES header and split it into 192-byte BD-TS
/// packets. `keyframe` drives the RAI bit on the first packet (video
/// only). The PES header and ES bytes are sliced in place — no second
/// full-frame copy.
fn write_pes_chain(
&mut self,
track: usize,
pid: u16,
pts_90k: u64,
is_video: bool,
keyframe: bool,
es_data: &[u8],
) -> io::Result<()> {
let pes_header = build_pes_header(pid, pts_90k, es_data.len());
// Logical PES packet = header bytes followed by es_data. It is
// indexed (and written) in place, without materializing the
// concatenation, to avoid a second full-frame copy on the hot path.
let pes_len = pes_header.len() + es_data.len();
let mut offset = 0;
let mut first = true;
while offset < pes_packet.len() {
let remaining = pes_packet.len() - offset;
while offset < pes_len {
let remaining = pes_len - offset;
// Invariant: TP_extra(4) + TS_header(4) + AF(af_bytes) + payload(payload_len) = 192,
// i.e. af_bytes + payload_len = TS_PAYLOAD (184).
@@ -165,8 +260,20 @@ impl<W: Write> TsMuxer<W> {
}
}
self.writer
.write_all(&pes_packet[offset..offset + payload_len])?;
// Write the payload span [offset, offset+payload_len), which may
// straddle the header/es_data boundary — emit each side in one
// write_all rather than copying the whole frame again.
let end = offset + payload_len;
let hdr_len = pes_header.len();
if offset < hdr_len {
let hdr_end = end.min(hdr_len);
self.writer.write_all(&pes_header[offset..hdr_end])?;
}
if end > hdr_len {
let es_start = offset.max(hdr_len) - hdr_len;
let es_end = end - hdr_len;
self.writer.write_all(&es_data[es_start..es_end])?;
}
offset += payload_len;
first = false;
@@ -175,6 +282,8 @@ impl<W: Write> TsMuxer<W> {
Ok(())
}
/// Flush the underlying writer. BD-TS needs no stream trailer, so this
/// only drains buffering; the muxer remains usable afterwards.
pub fn finish(&mut self) -> io::Result<()> {
self.writer.flush()
}
@@ -183,7 +292,7 @@ impl<W: Write> TsMuxer<W> {
/// Build a PES packet header for a BD stream.
fn build_pes_header(pid: u16, pts_90k: u64, data_len: usize) -> Vec<u8> {
// Determine stream_id from PID range
let stream_id: u8 = if (0x1011..=0x101F).contains(&pid) {
let stream_id: u8 = if is_video_pid(pid) {
0xE0 // video
} else {
0xBD // audio, PGS subtitle, or default (private stream 1)
@@ -198,7 +307,10 @@ fn build_pes_header(pid: u16, pts_90k: u64, data_len: usize) -> Vec<u8> {
header.push(0x01);
header.push(stream_id);
// PES packet length (0 = unbounded for video or if too large for u16)
// PES packet length. The unbounded form (0) is only spec-legal for
// video; `write_frame` splits oversized 0xBD access units so a private
// stream always fits a bounded u16 length here. The `> 65535` arm
// remains a defensive fallback for video only.
if stream_id == 0xE0 || pes_data_len > 65535 {
header.push(0x00);
header.push(0x00);
@@ -226,72 +338,6 @@ fn build_pes_header(pid: u16, pts_90k: u64, data_len: usize) -> Vec<u8> {
header
}
/// Extract NAL arrays from HEVCDecoderConfigurationRecord and convert to Annex B.
/// Returns VPS + SPS + PPS as Annex B NAL units (00 00 00 01 + NAL).
fn hvcc_to_annex_b(hvcc: &[u8]) -> Option<Vec<u8>> {
// HEVCDecoderConfigurationRecord: 22 bytes header, then NAL arrays
if hvcc.len() < 23 {
return None;
}
let num_arrays = hvcc[22] as usize;
let mut out = Vec::new();
let mut offset = 23;
for _ in 0..num_arrays {
if offset + 3 > hvcc.len() {
break;
}
// array: 1 byte (completeness + NAL type), 2 bytes (numNalus)
let _nal_type = hvcc[offset] & 0x3F;
let num_nalus = u16::from_be_bytes([hvcc[offset + 1], hvcc[offset + 2]]) as usize;
offset += 3;
for _ in 0..num_nalus {
if offset + 2 > hvcc.len() {
break;
}
let nal_len = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize;
offset += 2;
if offset + nal_len > hvcc.len() {
break;
}
out.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]);
out.extend_from_slice(&hvcc[offset..offset + nal_len]);
offset += nal_len;
}
}
if out.is_empty() { None } else { Some(out) }
}
/// Convert length-prefixed NALUs (4-byte BE length + NAL) to Annex B
/// (00 00 00 01 + NAL). Used for video elementary streams in TS.
fn length_prefixed_to_annex_b(data: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(data.len());
let mut offset = 0;
while offset + 4 <= data.len() {
let len = u32::from_be_bytes([
data[offset],
data[offset + 1],
data[offset + 2],
data[offset + 3],
]) as usize;
offset += 4;
if offset + len > data.len() {
break;
}
out.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]);
out.extend_from_slice(&data[offset..offset + len]);
offset += len;
}
// If data doesn't look like length-prefixed NALs (no valid parse),
// return original data unchanged — it may already be Annex B.
if out.is_empty() && !data.is_empty() {
return data.to_vec();
}
out
}
#[cfg(test)]
mod tests {
use super::*;
@@ -444,7 +490,7 @@ mod tests {
let mut sink: Vec<u8> = Vec::new();
{
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
mux.set_codec_private(0, hvcc);
mux.set_codec_private(0, hvcc).unwrap();
// Non-IDR before any IDR: should be dropped.
let p = fake_hevc_nal(1, 50);
mux.write_frame(0, 0, false, &p).unwrap();
@@ -477,6 +523,38 @@ mod tests {
);
}
#[test]
fn empty_data_keyframe_arms_params_so_later_frames_survive() {
// An empty-data keyframe must still arm params_written; otherwise
// every subsequent non-key frame would be dropped by the
// pre-keyframe guard and the track would emit no real frames.
let mut sink: Vec<u8> = Vec::new();
{
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
// Keyframe with empty payload (e.g. a frame whose NALs were
// all stripped upstream) — anchors the stream.
mux.write_frame(0, 0, true, &[]).unwrap();
// Now a real non-key frame; it must NOT be dropped.
let p = fake_hevc_nal(1, 80);
mux.write_frame(0, 41_000_000, false, &p).unwrap();
mux.finish().unwrap();
}
let packets = parse_bd_ts(&sink);
// The non-key frame's NAL body byte (0x02 = (1<<1)) must appear in
// a video payload — proof it wasn't dropped.
let video_bytes: Vec<u8> = packets
.iter()
.filter(|p| p.pid == VIDEO_PID)
.flat_map(|p| p.payload.clone())
.collect();
assert!(
video_bytes
.windows(4)
.any(|w| w == [0x00, 0x00, 0x00, 0x01]),
"later non-key frame must survive after an empty-data keyframe"
);
}
#[test]
fn non_key_before_first_keyframe_dropped() {
let mut sink: Vec<u8> = Vec::new();
@@ -493,4 +571,96 @@ mod tests {
"non-key before first keyframe must be dropped"
);
}
const AUDIO_PID: u16 = 0x1100;
/// Decode the 33-bit PTS from the first PUSI packet on `pid`. Assumes
/// the PES header carries PTS (flags 0x80 at PES byte 7).
fn first_pts_90k(packets: &[TsPacket], pid: u16) -> u64 {
let pkt = packets
.iter()
.find(|p| p.pid == pid && p.pusi)
.expect("PUSI packet present");
// PES payload starts the packet payload: 00 00 01 stream_id len len
// flags1 flags2 hdr_len then 5 PTS bytes.
let p = &pkt.payload;
let pts = &p[9..14];
((((pts[0] >> 1) & 0x07) as u64) << 30)
| ((pts[1] as u64) << 22)
| (((pts[2] >> 1) as u64) << 15)
| ((pts[3] as u64) << 7)
| ((pts[4] >> 1) as u64)
}
#[test]
fn av_offset_preserved_with_audio_before_first_video() {
// Audio at t=0 arrives BEFORE the first video keyframe at t=1s.
// The global base must be seeded from the VIDEO frame so the
// audio/video PTS offset is preserved (audio earlier ⇒ saturates to
// 0, video lands at +1s = 90000 ticks), not both collapsed to 0.
let mut sink: Vec<u8> = Vec::new();
{
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID, AUDIO_PID]);
// Audio frame first, at PTS 0.
mux.write_frame(1, 0, false, &[0x0B, 0x77, 0x00, 0x00])
.unwrap();
// Video keyframe at PTS 1s — seeds the base.
let idr = fake_hevc_nal(19, 100);
mux.write_frame(0, 1_000_000_000, true, &idr).unwrap();
mux.finish().unwrap();
}
let packets = parse_bd_ts(&sink);
let video_pts = first_pts_90k(&packets, VIDEO_PID);
let audio_pts = first_pts_90k(&packets, AUDIO_PID);
// Video keyframe is the base ⇒ its relative PTS is 0.
assert_eq!(video_pts, 0, "video keyframe seeds the base at t=0");
// Audio arrived 1s earlier ⇒ saturates to 0, NOT lifted past video.
assert_eq!(audio_pts, 0, "earlier audio saturates to 0");
assert!(
audio_pts <= video_pts,
"audio must not be pulled ahead of the video base"
);
}
#[test]
fn out_of_range_track_errors() {
let mut sink: Vec<u8> = Vec::new();
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
let err = mux.write_frame(5, 0, true, &[0xAA]).unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
let err2 = mux.set_codec_private(5, vec![0u8; 4]).unwrap_err();
assert_eq!(err2.kind(), std::io::ErrorKind::InvalidInput);
}
#[test]
fn oversized_bd_audio_pes_is_split_and_bounded() {
// A private_stream_1 (0xBD) audio frame larger than the bounded PES
// limit must be split into multiple PES, each with a non-zero
// PES_packet_length (never the unbounded 0 form, which is illegal
// for 0xBD).
let mut sink: Vec<u8> = Vec::new();
let big: Vec<u8> = (0..(MAX_BD_PES_PAYLOAD + 5000))
.map(|i| (i & 0xFF) as u8)
.collect();
{
let mut mux = TsMuxer::new(&mut sink, &[AUDIO_PID]);
mux.write_frame(0, 0, false, &big).unwrap();
mux.finish().unwrap();
}
let packets = parse_bd_ts(&sink);
let pusi: Vec<&TsPacket> = packets
.iter()
.filter(|p| p.pid == AUDIO_PID && p.pusi)
.collect();
assert!(
pusi.len() >= 2,
"oversized audio must span ≥2 PES, got {}",
pusi.len()
);
for p in pusi {
// PES length field at payload bytes [4..6] must be non-zero.
let len = u16::from_be_bytes([p.payload[4], p.payload[5]]);
assert_ne!(len, 0, "0xBD PES must carry a bounded length");
}
}
}