From 5287dd65ceaea90684c546c01d9294cdd6198b9d Mon Sep 17 00:00:00 2001 From: MattJackson <1085847+MattJackson@users.noreply.github.com> Date: Wed, 15 Apr 2026 16:22:28 +0000 Subject: [PATCH] Fix all PES pipeline audit findings (20 issues) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Critical: - C1: PES serialize validates track < 256 and data < 4GB - C2: PES deserialize caps frame size at 256MB (OOM protection) - C3: TsMuxer stuffing uses static buffer, no per-packet alloc - C4: PES length uses unbounded (0x0000) for audio >65535 bytes - C6: TsDemuxer validates AF length <= 183 Warning: - W1: parse_timestamp validates marker bits, returns Option - W2: PES header data_start clamped to data.len() - W3: TsMuxer PTS conversion uses saturating_mul, rejects negative - W4: AC3/DTS replace debug_assert with runtime bounds check - W6: MKV block_vint handles 3-4 byte VINTs - W7: meta.rs to_title() uses unwrap_or fallbacks instead of panic - W8: MKV reader skips frames for non-existent tracks - W9: DVD PTS uses higher-precision conversion (1e9/90000) - FMKV read_header caps JSON at 10MB - PAT section_len underflow guard Suggestion: - S2: TsMuxer uses static STUFF_FF buffer - S3: HEVC parser single-pass NAL scan (was duplicated) - S4: TsDemuxer caps remainder at one packet - S5: PTS 90kHz→ns uses round-to-nearest --- src/mux/codec/ac3.rs | 8 +++++-- src/mux/codec/dts.rs | 4 +++- src/mux/codec/h264.rs | 3 +++ src/mux/codec/hevc.rs | 52 +++++++++++++------------------------------ src/mux/codec/mod.rs | 6 +++-- src/mux/disc.rs | 2 +- src/mux/meta.rs | 26 +++++++++++++++------- src/mux/mkvstream.rs | 13 ++++++++++- src/mux/ts.rs | 48 ++++++++++++++++++++++++++++----------- src/mux/tsmux.rs | 21 +++++++++++------ src/pes.rs | 27 ++++++++++++++++++++-- 11 files changed, 136 insertions(+), 74 deletions(-) diff --git a/src/mux/codec/ac3.rs b/src/mux/codec/ac3.rs index edc06e4..3c4fe7d 100644 --- a/src/mux/codec/ac3.rs +++ b/src/mux/codec/ac3.rs @@ -99,7 +99,9 @@ fn find_ac3_sync(data: &[u8]) -> Option { /// bsid is at byte 5, bits 7..3. /// AC-3: bsid <= 10, E-AC-3: bsid >= 11 (typically 16). pub fn get_bsid(data: &[u8]) -> u8 { - debug_assert!(data.len() >= 6); + if data.len() < 6 { + return 0; + } (data[5] >> 3) & 0x1F } @@ -107,7 +109,9 @@ pub fn get_bsid(data: &[u8]) -> u8 { /// frmsiz is at bits [2:0] of byte 2 concatenated with byte 3. /// Frame size = (frmsiz + 1) * 2 bytes. pub fn eac3_frame_size(data: &[u8]) -> usize { - debug_assert!(data.len() >= 4); + if data.len() < 4 { + return 0; + } let frmsiz = ((data[2] as usize & 0x07) << 8) | (data[3] as usize); (frmsiz + 1) * 2 } diff --git a/src/mux/codec/dts.rs b/src/mux/codec/dts.rs index 51977c4..bd46e8c 100644 --- a/src/mux/codec/dts.rs +++ b/src/mux/codec/dts.rs @@ -82,7 +82,9 @@ pub fn find_dts_hd_ext_sync(data: &[u8]) -> Option { /// The size field is at bytes 6-8 of the extension: /// ((ext[6] & 0x1F) << 11) | (ext[7] << 3) | (ext[8] >> 5) + 1 pub fn dts_hd_ext_frame_size(ext: &[u8]) -> usize { - debug_assert!(ext.len() >= 9); + if ext.len() < 9 { + return 0; + } let raw = ((ext[6] as usize & 0x1F) << 11) | ((ext[7] as usize) << 3) | ((ext[8] as usize) >> 5); raw + 1 diff --git a/src/mux/codec/h264.rs b/src/mux/codec/h264.rs index fd43267..9b6a29f 100644 --- a/src/mux/codec/h264.rs +++ b/src/mux/codec/h264.rs @@ -180,6 +180,9 @@ pub fn find_start_code(data: &[u8], from: usize) -> Option { if data.len() < from + 3 { return None; } + // Range excludes last 2 bytes since we read 3 bytes at each position. + // data.len()-2 as exclusive upper bound means last checked index is data.len()-3, + // which accesses data[len-3], data[len-2], data[len-1] — all valid. (from..data.len() - 2).find(|&i| data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x01) } diff --git a/src/mux/codec/hevc.rs b/src/mux/codec/hevc.rs index 58c4cb9..2fff94b 100644 --- a/src/mux/codec/hevc.rs +++ b/src/mux/codec/hevc.rs @@ -52,8 +52,9 @@ impl CodecParser for HevcParser { 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 frame_data = Vec::new(); - // Scan NAL units + // Single-pass NAL scan: extract params, detect keyframes, build length-prefixed output let mut pos = 0; while let Some(sc_pos) = find_start_code(data, pos) { if let Some(nal_start) = skip_start_code(data, sc_pos) { @@ -68,45 +69,22 @@ impl CodecParser for HevcParser { let nal_type = (data[nal_start] >> 1) & 0x3F; match nal_type { - NAL_VPS => self.vps = Some(data[nal_start..end].to_vec()), - NAL_SPS => self.sps = Some(data[nal_start..end].to_vec()), - NAL_PPS => self.pps = Some(data[nal_start..end].to_vec()), + NAL_VPS => { self.vps = Some(data[nal_start..end].to_vec()); } + NAL_SPS => { self.sps = Some(data[nal_start..end].to_vec()); } + NAL_PPS => { self.pps = Some(data[nal_start..end].to_vec()); } + NAL_AUD => {} // Skip access unit delimiters t if (NAL_BLA_W_LP..=NAL_RSV_IRAP_VCL23).contains(&t) => { keyframe = true; + let nal = &data[nal_start..end]; + frame_data.extend_from_slice(&(nal.len() as u32).to_be_bytes()); + frame_data.extend_from_slice(nal); + } + _ => { + // All other NAL types (slices, SEI, DV RPU, etc.) pass through + let nal = &data[nal_start..end]; + frame_data.extend_from_slice(&(nal.len() as u32).to_be_bytes()); + frame_data.extend_from_slice(nal); } - _ => {} - } - } - pos = next; - } else { - break; - } - } - - // Convert Annex B to length-prefixed NALUs. - // Skip VPS/SPS/PPS/AUD — they're in codecPrivate. - let mut frame_data = Vec::new(); - let mut pos = 0; - while let Some(sc_pos) = find_start_code(&pes.data, pos) { - if let Some(nal_start) = skip_start_code(&pes.data, sc_pos) { - let next = find_start_code(&pes.data, nal_start).unwrap_or(pes.data.len()); - let mut end = next; - while end > nal_start && pes.data[end - 1] == 0x00 { - end -= 1; - } - - if nal_start < pes.data.len() { - let nal_type = (pes.data[nal_start] >> 1) & 0x3F; - // Skip parameter sets and AUD - if nal_type != NAL_VPS - && nal_type != NAL_SPS - && nal_type != NAL_PPS - && nal_type != NAL_AUD - { - let nal = &pes.data[nal_start..end]; - let len = nal.len() as u32; - frame_data.extend_from_slice(&len.to_be_bytes()); - frame_data.extend_from_slice(nal); } } pos = next; diff --git a/src/mux/codec/mod.rs b/src/mux/codec/mod.rs index bbf4b5f..cf8872b 100644 --- a/src/mux/codec/mod.rs +++ b/src/mux/codec/mod.rs @@ -31,9 +31,11 @@ pub struct Frame { pub data: Vec, } -/// Convert 90kHz PTS to nanoseconds. +/// Convert 90kHz PTS to nanoseconds (round to nearest). pub fn pts_to_ns(pts: i64) -> i64 { - pts * 100_000 / 9 + // pts * 1_000_000_000 / 90_000 = pts * 100_000 / 9 + // Add half-divisor for rounding: (pts * 100_000 + 4) / 9 + (pts * 100_000 + 4) / 9 } /// Trait for codec-specific elementary stream parsers. diff --git a/src/mux/disc.rs b/src/mux/disc.rs index 4e5b38c..1267d5f 100644 --- a/src/mux/disc.rs +++ b/src/mux/disc.rs @@ -401,7 +401,7 @@ impl crate::pes::Stream for DiscStream { _ => continue, }; if track < self.title.streams.len() { - let pts_ns = ps.pts.map(|p| (p as i64) * 100_000 / 9).unwrap_or(0); + let pts_ns = ps.pts.map(|p| (p as i64) * 1_000_000_000 / 90_000).unwrap_or(0); self.pending_frames.push_back(crate::pes::PesFrame { track, pts: pts_ns, diff --git a/src/mux/meta.rs b/src/mux/meta.rs index f199c0e..2c1c187 100644 --- a/src/mux/meta.rs +++ b/src/mux/meta.rs @@ -150,10 +150,10 @@ impl M2tsMeta { codec_private: _, } => Stream::Video(VideoStream { pid: *pid, - codec: codec.parse().unwrap(), - resolution: resolution.parse().unwrap(), - frame_rate: frame_rate.parse().unwrap(), - hdr: hdr.parse().unwrap(), + codec: codec.parse().unwrap_or(crate::disc::Codec::Unknown(0)), + resolution: resolution.parse().unwrap_or(crate::disc::Resolution::Unknown), + frame_rate: frame_rate.parse().unwrap_or(crate::disc::FrameRate::Unknown), + hdr: hdr.parse().unwrap_or(crate::disc::HdrFormat::Sdr), color_space: ColorSpace::Bt709, secondary: *secondary, label: label.clone(), @@ -168,10 +168,10 @@ impl M2tsMeta { secondary, } => Stream::Audio(AudioStream { pid: *pid, - codec: codec.parse().unwrap(), - channels: channels.parse().unwrap(), + codec: codec.parse().unwrap_or(crate::disc::Codec::Unknown(0)), + channels: channels.parse().unwrap_or(crate::disc::AudioChannels::Unknown), language: language.clone(), - sample_rate: sample_rate.parse().unwrap(), + sample_rate: sample_rate.parse().unwrap_or(crate::disc::SampleRate::Unknown), secondary: *secondary, label: label.clone(), }), @@ -182,7 +182,7 @@ impl M2tsMeta { forced, } => Stream::Subtitle(SubtitleStream { pid: *pid, - codec: codec.parse().unwrap(), + codec: codec.parse().unwrap_or(crate::disc::Codec::Unknown(0)), language: language.clone(), forced: *forced, codec_data: None, @@ -256,9 +256,14 @@ pub fn read_header(r: &mut R) -> io::Result> { return Ok(None); } + const MAX_JSON_SIZE: usize = 10 * 1024 * 1024; // 10 MB + let mut len_buf = [0u8; 4]; r.read_exact(&mut len_buf)?; let json_len = u32::from_be_bytes(len_buf) as usize; + if json_len > MAX_JSON_SIZE { + return Err(io::Error::new(io::ErrorKind::InvalidData, "FMKV JSON too large")); + } let mut json_buf = vec![0u8; json_len]; r.read_exact(&mut json_buf)?; @@ -287,9 +292,14 @@ pub fn read_header_from_stream(r: &mut impl Read) -> io::Result return Ok(None); } + const MAX_JSON_SIZE: usize = 10 * 1024 * 1024; + let mut len_buf = [0u8; 4]; r.read_exact(&mut len_buf)?; let json_len = u32::from_be_bytes(len_buf) as usize; + if json_len > MAX_JSON_SIZE { + return Err(io::Error::new(io::ErrorKind::InvalidData, "FMKV JSON too large")); + } let mut json_buf = vec![0u8; json_len]; r.read_exact(&mut json_buf)?; diff --git a/src/mux/mkvstream.rs b/src/mux/mkvstream.rs index 90ac114..b0787b1 100644 --- a/src/mux/mkvstream.rs +++ b/src/mux/mkvstream.rs @@ -183,6 +183,11 @@ impl crate::pes::Stream for MkvStream { let pts_ms = rs.cluster_ts_ms + rel_ts as i64; let track_idx = (track as usize).saturating_sub(1); // MKV tracks are 1-based + // Skip blocks for non-existent tracks + if track_idx >= self.disc_title.streams.len() { + continue; + } + return Ok(Some(crate::pes::PesFrame { track: track_idx, pts: pts_ms * 1_000_000, // ms → ns @@ -689,7 +694,13 @@ fn block_vint(d: &[u8]) -> (u64, usize) { if d[0] & 0x40 != 0 && d.len() >= 2 { return ((((d[0] & 0x3F) as u64) << 8) | d[1] as u64, 2); } - (0, 1) + 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) // Unsupported 5+ byte VINT — treat as track 0 } /// Convert HEVCDecoderConfigurationRecord (hvcC) to Annex B NAL units. diff --git a/src/mux/ts.rs b/src/mux/ts.rs index e159457..382427b 100644 --- a/src/mux/ts.rs +++ b/src/mux/ts.rs @@ -177,6 +177,9 @@ impl TsDemuxer { let payload_start = if adaptation == 0x03 || adaptation == 0x02 { // Adaptation field present let af_len = ts[4] as usize; + if af_len > 183 { + continue; // Malformed: AF length exceeds TS payload + } 5 + af_len } else { 4 @@ -208,9 +211,15 @@ impl TsDemuxer { } } - // Save leftover bytes for next call + // Save leftover bytes for next call (cap at one packet to prevent unbounded growth) if offset < work.len() { - self.remainder.extend_from_slice(&work[offset..]); + let leftover = &work[offset..]; + if leftover.len() < BD_TS_PACKET_SIZE { + self.remainder.extend_from_slice(leftover); + } else { + // More than one full packet leftover — something is wrong, discard + self.remainder.clear(); + } } completed @@ -257,30 +266,38 @@ fn parse_pes_header(data: &[u8]) -> (Option, Option, usize) { let pts_dts_flags = (data[7] >> 6) & 0x03; let header_data_len = data[8] as usize; - let data_start = 9 + header_data_len; + let data_start = (9 + header_data_len).min(data.len()); let mut pts = None; let mut dts = None; - if pts_dts_flags >= 2 && data.len() >= 14 { - pts = Some(parse_timestamp(&data[9..14])); + if pts_dts_flags >= 2 && header_data_len >= 5 && data.len() >= 14 { + pts = parse_timestamp(&data[9..14]); } - if pts_dts_flags == 3 && data.len() >= 19 { - dts = Some(parse_timestamp(&data[14..19])); + if pts_dts_flags == 3 && header_data_len >= 10 && data.len() >= 19 { + dts = parse_timestamp(&data[14..19]); } (pts, dts, data_start) } /// Parse a 5-byte PTS/DTS timestamp (33 bits in 90kHz). -fn parse_timestamp(data: &[u8]) -> i64 { +/// Validates marker bits per MPEG-2 spec. Returns None on invalid encoding. +fn parse_timestamp(data: &[u8]) -> Option { + if data.len() < 5 { + return None; + } + // Validate marker bits: byte 2 bit 0 and byte 4 bit 0 must be 1 + if (data[2] & 0x01) == 0 || (data[4] & 0x01) == 0 { + return None; + } let b0 = data[0] as i64; let b1 = data[1] as i64; let b2 = data[2] as i64; let b3 = data[3] as i64; let b4 = data[4] as i64; - ((b0 >> 1) & 0x07) << 30 | b1 << 22 | (b2 >> 1) << 15 | b3 << 7 | b4 >> 1 + Some(((b0 >> 1) & 0x07) << 30 | b1 << 22 | (b2 >> 1) << 15 | b3 << 7 | b4 >> 1) } // ============================================================ @@ -312,6 +329,7 @@ pub fn scan_streams(data: &[u8]) -> Option> { let section_len = (((data[pat_start + 1] & 0x0F) as usize) << 8) | data[pat_start + 2] as usize; let entries_start = pat_start + 8; + if section_len < 4 { offset += BD_TS_PACKET_SIZE; continue; } let entries_end = pat_start + 3 + section_len - 4; let mut e = entries_start; while e + 4 <= data.len() && e < entries_end { @@ -507,7 +525,7 @@ pub fn scan_first_pts(data: &[u8], target_pid: u16) -> Option { if payload.len() >= 14 && payload[0] == 0 && payload[1] == 0 && payload[2] == 1 { let pts_dts_flags = (payload[7] >> 6) & 0x03; if pts_dts_flags >= 2 { - return Some(parse_timestamp(&payload[9..14])); + return parse_timestamp(&payload[9..14]); } } } @@ -537,7 +555,7 @@ pub fn scan_last_pts(data: &[u8], target_pid: u16) -> Option { if payload.len() >= 14 && payload[0] == 0 && payload[1] == 0 && payload[2] == 1 { let pts_dts_flags = (payload[7] >> 6) & 0x03; if pts_dts_flags >= 2 { - last_pts = Some(parse_timestamp(&payload[9..14])); + last_pts = parse_timestamp(&payload[9..14]); } } } @@ -589,14 +607,18 @@ mod tests { fn test_parse_timestamp() { // Example: PTS = 0 → encoded as 21 00 01 00 01 let data = [0x21, 0x00, 0x01, 0x00, 0x01]; - assert_eq!(parse_timestamp(&data), 0); + assert_eq!(parse_timestamp(&data), Some(0)); // Example: PTS = 90000 (1 second at 90kHz) // Manual encoding: 33 bits = 0x00015F90 // This is just a sanity check that the parser doesn't crash let data2 = [0x21, 0x00, 0x07, 0xE9, 0x01]; // approximate let pts = parse_timestamp(&data2); - assert!(pts >= 0); + assert!(pts.is_some() && pts.unwrap() >= 0); + + // Invalid marker bits → returns None + let bad = [0x00, 0x00, 0x00, 0x00, 0x00]; // marker bits wrong + assert_eq!(parse_timestamp(&bad), None); } #[test] diff --git a/src/mux/tsmux.rs b/src/mux/tsmux.rs index 5aa53ed..8d1fee6 100644 --- a/src/mux/tsmux.rs +++ b/src/mux/tsmux.rs @@ -72,7 +72,11 @@ impl TsMuxer { }; // Build PES packet: header + data - let pts_90k = (pts_ns * 9 / 100_000) as u64; + 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(); @@ -108,14 +112,17 @@ impl TsMuxer { self.writer.write_all(&tp_extra)?; self.writer.write_all(&ts_header)?; + // Write adaptation field: length byte + flags byte + 0xFF padding + // stuff_len == 1: AF length = 0 (just the length byte, no flags) + // stuff_len >= 2: AF length = stuff_len-1, flags = 0, rest 0xFF + static STUFF_FF: [u8; 184] = [0xFF; 184]; if stuff_len == 1 { self.writer.write_all(&[0u8])?; // adaptation_field_length = 0 } else { - self.writer.write_all(&[(stuff_len - 1) as u8])?; // length + self.writer.write_all(&[(stuff_len - 1) as u8])?; // AF length self.writer.write_all(&[0u8])?; // flags if stuff_len > 2 { - let padding = vec![0xFF; stuff_len - 2]; - self.writer.write_all(&padding)?; + self.writer.write_all(&STUFF_FF[..stuff_len - 2])?; } } self.writer.write_all(&pes_packet[offset..offset + payload_len])?; @@ -159,12 +166,12 @@ fn build_pes_header(pid: u16, pts_90k: u64, data_len: usize) -> Vec { header.push(0x01); header.push(stream_id); - // PES packet length (0 = unbounded for video) - if stream_id == 0xE0 { + // PES packet length (0 = unbounded for video or if too large for u16) + if stream_id == 0xE0 || pes_data_len > 65535 { header.push(0x00); header.push(0x00); } else { - let len = (pes_data_len & 0xFFFF) as u16; + let len = pes_data_len as u16; header.push((len >> 8) as u8); header.push(len as u8); } diff --git a/src/pes.rs b/src/pes.rs index c7963fc..3d83733 100644 --- a/src/pes.rs +++ b/src/pes.rs @@ -22,6 +22,18 @@ pub struct PesFrame { impl PesFrame { /// Serialize to bytes: track(1) | pts(8) | keyframe(1) | len(4) | data pub fn serialize(&self, w: &mut dyn std::io::Write) -> std::io::Result<()> { + if self.track > 255 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "track index exceeds 255", + )); + } + if self.data.len() > u32::MAX as usize { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "frame data exceeds 4 GB", + )); + } w.write_all(&[self.track as u8])?; w.write_all(&self.pts.to_le_bytes())?; w.write_all(&[if self.keyframe { 1 } else { 0 }])?; @@ -31,6 +43,8 @@ impl PesFrame { /// Deserialize from bytes. Returns None at EOF. pub fn deserialize(r: &mut dyn std::io::Read) -> std::io::Result> { + const MAX_FRAME_SIZE: usize = 256 * 1024 * 1024; // 256 MB + let mut header = [0u8; 14]; // 1 + 8 + 1 + 4 match r.read_exact(&mut header) { Ok(_) => {} @@ -38,9 +52,18 @@ impl PesFrame { Err(e) => return Err(e), } let track = header[0] as usize; - let pts = i64::from_le_bytes(header[1..9].try_into().unwrap()); + let pts = i64::from_le_bytes([ + header[1], header[2], header[3], header[4], + header[5], header[6], header[7], header[8], + ]); let keyframe = header[9] != 0; - let len = u32::from_le_bytes(header[10..14].try_into().unwrap()) as usize; + let len = u32::from_le_bytes([header[10], header[11], header[12], header[13]]) as usize; + if len > MAX_FRAME_SIZE { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("frame size {} exceeds maximum {}", len, MAX_FRAME_SIZE), + )); + } let mut data = vec![0u8; len]; r.read_exact(&mut data)?; Ok(Some(Self { track, pts, keyframe, data }))