From 81b7e3ab4e78b2093e6c8cfc1241b5852e9188d1 Mon Sep 17 00:00:00 2001 From: Matthew Jackson <1085847+MattJackson@users.noreply.github.com> Date: Wed, 30 Sep 2026 15:50:09 -0700 Subject: [PATCH] ts: robust transport-stream demux Skips duplicate packets, gathers PES headers that span packets, resyncs after byte slips and across zero-filled gaps, prefers a current CRC-valid PAT/PMT, and drops the open PES on transport_error_indicator. M2TS output remaps PIDs into the BD ranges. --- src/mux/decode_ts.rs | 48 +++- src/mux/decode_ts/tests.rs | 35 +++ src/mux/demux_thread.rs | 4 +- src/mux/m2ts.rs | 194 ++++++++++++- src/mux/ts.rs | 544 ++++++++++++++++++++++++++++++------- src/mux/tsmux.rs | 62 ++++- 6 files changed, 757 insertions(+), 130 deletions(-) diff --git a/src/mux/decode_ts.rs b/src/mux/decode_ts.rs index 63688927..e846f8dc 100644 --- a/src/mux/decode_ts.rs +++ b/src/mux/decode_ts.rs @@ -401,17 +401,19 @@ fn vc1_param(bdu: &[u8]) -> Option { (*bdu.first()? == 0x0F).then_some(1) } -// NAL units of an IR frame: length-prefixed (`nal_len` octets), or Annex B when the frame -// already starts with a start code (the probe `length_prefixed_to_annex_b` uses). +// NAL units of an IR frame: length-prefixed (`nal_len` octets) as the writer reads it; Annex B +// only when it starts with a start code and the prefixes do not exactly tile the frame. fn nals(es: &[u8], nal_len: usize) -> Nals<'_> { - if es.starts_with(&[0, 0, 1]) || es.starts_with(&[0, 0, 0, 1]) { + let prefixed = LengthPrefixed { + es, + nal_len, + pos: 0, + }; + let start_code = es.starts_with(&[0, 0, 1]) || es.starts_with(&[0, 0, 0, 1]); + if start_code && !prefixed.tiles() { Nals::AnnexB(AnnexB::new(es)) } else { - Nals::Prefixed(LengthPrefixed { - es, - nal_len, - pos: 0, - }) + Nals::Prefixed(prefixed) } } @@ -436,14 +438,34 @@ struct LengthPrefixed<'a> { pos: usize, } -impl<'a> Iterator for LengthPrefixed<'a> { - type Item = &'a [u8]; - fn next(&mut self) -> Option<&'a [u8]> { - let size = if (1..=4).contains(&self.nal_len) { +impl LengthPrefixed<'_> { + fn size(&self) -> usize { + if (1..=4).contains(&self.nal_len) { self.nal_len } else { 4 - }; + } + } + + // Whether the length prefixes land exactly on the end of the frame. + fn tiles(&self) -> bool { + let size = self.size(); + let mut pos = 0usize; + while pos < self.es.len() { + let Some(head) = self.es.get(pos..pos + size) else { + return false; + }; + let len = head.iter().fold(0usize, |a, &b| (a << 8) | b as usize); + pos = pos.saturating_add(size).saturating_add(len); + } + pos == self.es.len() + } +} + +impl<'a> Iterator for LengthPrefixed<'a> { + type Item = &'a [u8]; + fn next(&mut self) -> Option<&'a [u8]> { + let size = self.size(); loop { let head = self.es.get(self.pos..self.pos + size)?; let len = head.iter().fold(0usize, |a, &b| (a << 8) | b as usize); diff --git a/src/mux/decode_ts/tests.rs b/src/mux/decode_ts/tests.rs index fdf11b8a..b184402d 100644 --- a/src/mux/decode_ts/tests.rs +++ b/src/mux/decode_ts/tests.rs @@ -628,6 +628,41 @@ fn hevc_b_pyramid_uses_the_sps_reorder_depth() { assert_eq!(dts, vec![f(-2), f(-1), f(0), None, f(2)]); } +// A 4-octet prefix of 1 or 256..=511 starts `00 00 00 01` / `00 00 01`: the frame is +// still length-prefixed (the writer never probes), so the in-band SPS must be found. +#[test] +fn length_prefixed_frame_whose_prefix_looks_like_a_start_code_is_not_annex_b() { + for first in [vec![0x4C], { + let mut fd = vec![0x4C, 0x01]; + fd.resize(300, 0xFF); + fd + }] { + let mut hvcc = vec![0u8; 23]; + hvcc[21] = 0x03; + let mut d = DtsDeriver::for_codec(Codec::Hevc, Some(&hvcc)); + let sps = hevc_sps_vui(2, Some((1, 24)), false); + let frames: Vec<_> = [0i64, 4, 2, 1, 3] + .into_iter() + .enumerate() + .map(|(i, o)| { + let mut au = vec![first.clone(), hevc_slice(1)]; + if i == 0 { + au.insert(1, sps.clone()); + } + (B + o * F, length_prefixed(&au)) + }) + .collect(); + let f = |n: i64| Some(B + n * F); + let want = vec![f(-2), f(-1), f(0), None, f(2)]; + assert_eq!( + run(&mut d, &frames), + want, + "first NAL {} bytes", + first.len() + ); + } +} + #[test] fn hevc_vui_timing_is_reached_through_scaling_lists_pcm_and_rps() { for scaling in [false, true] { diff --git a/src/mux/demux_thread.rs b/src/mux/demux_thread.rs index 1bb788d1..b5811d99 100644 --- a/src/mux/demux_thread.rs +++ b/src/mux/demux_thread.rs @@ -404,7 +404,9 @@ mod tests { let (_dt, rx) = DemuxThread::spawn_zero_copy(pf_rx, rc_tx, (), None, Some(ts), None).unwrap(); pf_tx.send(Ok(bdts_pes_packet(pid, &[0x01]))).unwrap(); - pf_tx.send(Ok(bdts_pes_packet(pid, &[0x02]))).unwrap(); + let mut second = bdts_pes_packet(pid, &[0x02]); + second[7] |= 1; // next CC: a same-CC packet is a duplicate + pf_tx.send(Ok(second)).unwrap(); drop(pf_tx); let batches = collect_batches(&rx, Duration::from_secs(5)); diff --git a/src/mux/m2ts.rs b/src/mux/m2ts.rs index a0e12fd8..c3e67513 100644 --- a/src/mux/m2ts.rs +++ b/src/mux/m2ts.rs @@ -19,6 +19,33 @@ pub(crate) fn lpcm_bd_header(a: &crate::disc::AudioStream, cp: Option<&[u8]>) -> super::codec::lpcm::bd_header(a.channels.count(), a.sample_rate.hz() as u32, src) } +fn stream_pid(s: &DiscStream) -> u16 { + match s { + DiscStream::Video(v) => v.pid, + DiscStream::Audio(a) => a.pid, + DiscStream::Subtitle(s) => s.pid, + } +} + +// tsmux picks stream_id and NAL handling by PID, so video must sit in its video range and +// nothing else may; DVD ids (0xE0, 0xBD80) and the null PID are not carriable at all. +fn pid_fits(s: &DiscStream, pid: u16) -> bool { + let video = matches!(s, DiscStream::Video(_)); + video == super::tsmux::is_video_pid(pid) && (0x0010..=0x1FFE).contains(&pid) +} + +// Lowest free carriable PID for `s`, searched from its BD base. +fn free_pid(s: &DiscStream, used: &[bool]) -> Option { + let base = match s { + DiscStream::Video(_) => *super::tsmux::VIDEO_PID_RANGE.start(), + DiscStream::Audio(_) => 0x1100, + DiscStream::Subtitle(_) => 0x1200, + }; + (base..=0x1FFE) + .chain(0x0010..base) + .find(|&p| pid_fits(s, p) && !used[p as usize]) +} + /// BD transport stream write sink with embedded FMKV metadata /// header. pub struct M2tsStream { @@ -41,6 +68,21 @@ impl M2tsStream { out.codec_privates.clear(); let mut route = Vec::with_capacity(title.streams.len()); let excluded = super::ps::UnstoredExtensions::new(title, "M2TS"); + // Keep each carriable source PID (first use wins); the rest are remapped below and + // the header records the PID actually written. + let mut used = vec![false; 0x2000]; + let keep: Vec = title + .streams + .iter() + .map(|s| { + let pid = stream_pid(s); + let k = pid_fits(s, pid) && !used[pid as usize]; + if k { + used[pid as usize] = true; + } + k + }) + .collect(); for (i, s) in title.streams.iter().enumerate() { let mut cp = title.codec_privates.get(i).cloned().flatten(); // No descriptor binds a 13818-3 extension PID to its base, so a player could not @@ -67,8 +109,22 @@ impl M2tsStream { } _ => None, }; + let mut s = s.clone(); + if !keep[i] { + let Some(pid) = free_pid(&s, &used) else { + tracing::warn!(target: "mux", track = i, "no free BD-TS PID; track omitted from M2TS"); + route.push(None); + continue; + }; + used[pid as usize] = true; + match &mut s { + DiscStream::Video(v) => v.pid = pid, + DiscStream::Audio(a) => a.pid = pid, + DiscStream::Subtitle(t) => t.pid = pid, + } + } route.push(Some((out.streams.len(), lpcm))); - out.streams.push(s.clone()); + out.streams.push(s); out.codec_privates.push(cp); } // Write FMKV header unconditionally: skipping it for a zero-stream title @@ -76,15 +132,7 @@ impl M2tsStream { // read-back (read_header → Ok(None) → PMT fallback). Empty array is valid. let m = meta::M2tsMeta::from_title(&out); meta::write_header(&mut writer, &m)?; - let pids: Vec = out - .streams - .iter() - .map(|s| match s { - DiscStream::Video(v) => v.pid, - DiscStream::Audio(a) => a.pid, - DiscStream::Subtitle(s) => s.pid, - }) - .collect(); + let pids: Vec = out.streams.iter().map(stream_pid).collect(); let boxed: Box = Box::new(writer); let mut muxer = super::tsmux::TsMuxer::new(boxed, &pids); // Declaring the codec decides both ES framing (HEVC/H.264 are length- @@ -163,6 +211,7 @@ impl crate::pes::Stream for M2tsStream { self.muxer.finish() } + // The source title: its PIDs, not the BD-range PIDs the file and FMKV header carry. fn info(&self) -> &crate::disc::DiscTitle { &self.disc_title } @@ -783,6 +832,131 @@ mod tests { ); } + // Write one frame per track of `title` and return (header pids, BD-TS bytes). + fn mux_one_frame_each(title: &DiscTitle, data: &[Vec]) -> (Vec, Vec) { + let shared = std::sync::Arc::new(std::sync::Mutex::new(Vec::::new())); + let mut stream = M2tsStream::create(SharedSink(shared.clone()), title).unwrap(); + for (track, d) in data.iter().enumerate() { + stream.write(&frame(track, 0, true, d.clone())).unwrap(); + } + stream.finish().unwrap(); + drop(stream); + let buf = shared.lock().unwrap().clone(); + let (meta, ts) = ts_after_header(&buf); + let pids = meta + .to_title() + .streams + .iter() + .map(|s| match s { + DiscStream::Video(v) => v.pid, + DiscStream::Audio(a) => a.pid, + DiscStream::Subtitle(s) => s.pid, + }) + .collect(); + (pids, ts) + } + + fn ac3_audio(pid: u16) -> DiscStream { + DiscStream::Audio(crate::disc::AudioStream { + pid, + codec: Codec::Ac3, + channels: crate::disc::AudioChannels::Stereo, + language: "eng".into(), + sample_rate: crate::disc::SampleRate::S48, + secondary: false, + purpose: crate::disc::LabelPurpose::Normal, + label: String::new(), + }) + } + + // DVD ids (video 0xE0, audio 0xBD80) are not 13-bit TS PIDs: the header must name + // the PID the packets actually carry, and video must use the video path. + #[test] + fn dvd_stream_ids_get_bd_pids_that_the_header_records() { + let mut title = make_title(); + if let DiscStream::Video(v) = &mut title.streams[0] { + v.pid = 0xE0; + v.codec = Codec::H264; + } + title.streams.push(ac3_audio(0xBD80)); + title.codec_privates = vec![None, None]; + let audio = vec![0x0B, 0x77, 1, 2, 3, 4, 5, 6]; + let (pids, ts) = mux_one_frame_each(&title, &[fake_idr_pes_data(), audio.clone()]); + assert!(pids.iter().all(|&p| p <= 0x1FFF), "13-bit PIDs: {pids:x?}"); + let mut demux = crate::mux::ts::TsDemuxer::new(&pids); + let mut pes = demux.feed(&ts); + pes.extend(demux.flush()); + let got = |pid| pes.iter().find(|p| p.pid == pid).map(|p| p.data.clone()); + assert_eq!(got(pids[1]), Some(audio), "audio found on its header PID"); + let v = got(pids[0]).expect("video found on its header PID"); + assert!(v.starts_with(&[0, 0, 0, 1]), "video converted to Annex B"); + assert!( + pes_times(&ts, pids[0]).len() == 1, + "video PES on the header PID" + ); + } + + // An audio track on a video-range PID (MKV whose track 1 is audio) must pass verbatim. + #[test] + fn audio_on_a_video_range_pid_is_not_treated_as_video() { + let mut title = make_title(); + title.streams = vec![ac3_audio(0x1011)]; + title.codec_privates = vec![None]; + let audio = vec![0, 0, 0, 4, 0x0B, 0x77, 9, 9]; + let (pids, ts) = mux_one_frame_each(&title, std::slice::from_ref(&audio)); + let mut demux = crate::mux::ts::TsDemuxer::new(&pids); + let mut pes = demux.feed(&ts); + pes.extend(demux.flush()); + assert_eq!(pes.len(), 1); + assert_eq!(pes[0].data, audio, "audio bytes pass through unconverted"); + } + + // BD PiP (0x1B00) and MKV non-first video (0x1100+) move into the video range; BD PIDs + // already in their range are kept, and each track reads back on its header PID. + #[test] + fn out_of_range_video_moves_into_the_video_range_and_bd_pids_stay() { + let mut title = make_title(); + let DiscStream::Video(v) = title.streams[0].clone() else { + unreachable!() + }; + for pid in [0x1B00, 0x1100] { + title + .streams + .push(DiscStream::Video(VideoStream { pid, ..v.clone() })); + } + title.streams.push(ac3_audio(0x1101)); + title + .streams + .push(DiscStream::Subtitle(crate::disc::SubtitleStream { + pid: 0x1200, + codec: Codec::Pgs, + language: "eng".into(), + forced: false, + qualifier: crate::disc::LabelQualifier::None, + codec_data: None, + })); + title.codec_privates = vec![None; 5]; + let audio = vec![0x0B, 0x77, 1, 2]; + let sub = vec![0x16, 0, 0]; + let v = fake_idr_pes_data(); + let data = [v.clone(), v.clone(), v, audio.clone(), sub.clone()]; + let (pids, ts) = mux_one_frame_each(&title, &data); + assert_eq!(pids, [0x1011, 0x1012, 0x1013, 0x1101, 0x1200]); + let mut demux = crate::mux::ts::TsDemuxer::new(&pids); + let mut pes = demux.feed(&ts); + pes.extend(demux.flush()); + let got = |pid| pes.iter().find(|p| p.pid == pid).map(|p| p.data.clone()); + for &pid in &pids[..3] { + let v = got(pid).expect("video found on its header PID"); + assert!( + v.starts_with(&[0, 0, 0, 1]), + "video {pid:#x} converted to Annex B" + ); + } + assert_eq!(got(0x1101), Some(audio)); + assert_eq!(got(0x1200), Some(sub)); + } + #[test] fn mvc_dependent_view_matches_the_base_dts() { mvc_dependent_dts(false); diff --git a/src/mux/ts.rs b/src/mux/ts.rs index f3a255a7..c98cad42 100644 --- a/src/mux/ts.rs +++ b/src/mux/ts.rs @@ -64,6 +64,9 @@ struct PesAssembler { /// HEVC/H264 that reads as a spurious start code / corrupt slice /// payload. Tracks how many header bytes remain across packets. header_remaining: usize, + /// Payload bytes since the PUSI while the PES header (fixed part or PTS/DTS) is still + /// incomplete: h222 lets it span packets. Empty otherwise. + head: Vec, /// 4-bit continuity_counter of the last payload-bearing TS packet seen /// on this PID. A non-PUSI continuation whose CC is not `(prev + 1) & 0xf` /// — or whose adaptation field flags a discontinuity — means one or more @@ -71,6 +74,10 @@ struct PesAssembler { /// partial PES would inject corrupt bytes. The partial PES is dropped and /// the assembler resyncs on the next PUSI. `None` until the first packet. last_cc: Option, + /// PUSI bit and payload of that packet: a same-CC packet is a duplicate (h222 2.4.3.3) + /// only if it repeats them; otherwise it is new data after a gap (e.g. a clip join). + last_pusi: bool, + last_payload: Vec, /// Absolute source byte offset of the in-progress PES's first byte (the /// PUSI packet that began it), or `None` when no source base is threaded. /// Stamped at PES start, emitted on the completed packet — provenance is @@ -110,7 +117,10 @@ impl PesAssembler { dts: None, active: false, header_remaining: 0, + head: Vec::new(), last_cc: None, + last_pusi: false, + last_payload: Vec::new(), pes_source: None, pending_discontinuity: false, } @@ -164,16 +174,56 @@ impl PesAssembler { self.header_remaining = 0; // A dropped partial PES is a gap in the elementary stream — flag // it so the NEXT completed PES carries a discontinuity, matching - // every other partial-drop path in this file (lines 395/479/504). + // every other partial-drop path in this file (`drop_partial`). self.pending_discontinuity = true; return; } + // Grow by doubling but never past `cap`, so the allocation honours the share too. + let need = self.buffer.len() + data.len(); + if need > self.buffer.capacity() { + let target = need.max(self.buffer.capacity() * 2).min(self.cap); + self.buffer.reserve_exact(target - self.buffer.len()); + } self.buffer.extend_from_slice(data); } } + // Drop the open PES (its data has a hole) and flag the next one; resyncs on the next PUSI. + fn drop_partial(&mut self) { + self.buffer.clear(); + self.active = false; + self.header_remaining = 0; + self.head.clear(); + self.pending_discontinuity = true; + } + + // Take PUSI payload bytes into the PES header; once it and its PTS/DTS are complete, set + // the timestamps and pass the rest on as ES. A payload that is not a PES start opens no + // PES (its continuations would form a headless frame) and flags the loss. + fn take_head(&mut self, bytes: &[u8]) { + self.head.extend_from_slice(bytes); + let Some((pts, dts, header_len)) = pes_header_complete(&self.head) else { + return; + }; + let head = std::mem::take(&mut self.head); + if header_len == 0 { + self.drop_partial(); + } else { + self.pts = pts; + self.dts = dts; + if header_len < head.len() { + self.push(&head[header_len..]); + } else { + self.header_remaining = header_len - head.len(); + } + } + self.head = head; + self.head.clear(); + } + /// Flush remaining data as a PES packet. fn flush(&mut self) -> Option { + self.head.clear(); if self.active && !self.buffer.is_empty() { self.active = false; let discontinuity = self.pending_discontinuity; @@ -205,6 +255,8 @@ pub struct TsDemuxer { /// True once a caller has threaded a source base via [`Self::feed_at`]. Until /// then no `SourcePos` is stamped (keeps existing callers byte-identical). has_base: bool, + /// Set once sync loss has been logged, so a damaged input warns only once. + sync_lost_logged: bool, } impl TsDemuxer { @@ -237,6 +289,7 @@ impl TsDemuxer { remainder: Vec::new(), feed_base: 0, has_base: false, + sync_lost_logged: false, } } @@ -308,6 +361,12 @@ impl TsDemuxer { // Aligned-packets fast path — reads directly out of `data`. while offset + BD_SOURCE_PACKET_BYTES <= data.len() { + if data[offset + 4] != SYNC_BYTE { + // No sync byte: a zero-filled gap or a byte slip. Packets lost in a slip + // surface as CC gaps on their PIDs. + offset = self.resync(data, offset); + continue; + } let packet = &data[offset..offset + BD_SOURCE_PACKET_BYTES]; let src = self.pkt_source(offset); offset += BD_SOURCE_PACKET_BYTES; @@ -333,6 +392,30 @@ impl TsDemuxer { completed } + // Next packet start after the unsynced `offset`: the next grid slot if it has a sync byte, + // else the first grid slot with one or an off-grid point confirmed 192 bytes on (a slip). + // With neither, the last grid slot, so alignment carries over. + fn resync(&mut self, data: &[u8], offset: usize) -> usize { + let step = BD_SOURCE_PACKET_BYTES; + if data.get(offset + step + 4) == Some(&SYNC_BYTE) { + return offset + step; + } + for p in offset + 1..data.len().saturating_sub(4) { + if (p - offset).is_multiple_of(step) { + if data[p + 4] == SYNC_BYTE { + return p; + } + } else if data[p + 4] == SYNC_BYTE && data.get(p + step + 4) == Some(&SYNC_BYTE) { + if !self.sync_lost_logged { + self.sync_lost_logged = true; + tracing::warn!(target: "mux", offset = p, "bd-ts: packet sync slipped; resynced"); + } + return p; + } + } + offset + (data.len() - offset) / step * step + } + // Demux a single 192-byte BD-TS packet (4-byte TP_extra_header + 188-byte // TS). Routes payload bytes into the per-PID `PesAssembler`; completed PES // packets are pushed onto `completed` so the caller's alloc amortises. @@ -364,10 +447,7 @@ impl TsDemuxer { // A concealed unit may have dropped packets from any open PES, leaving a // hole mid-access-unit. Drop it like a mid-PES continuity break and flag // pending so the next completed PES resyncs (mirrors the non-PUSI cc_gap path). - a.buffer.clear(); - a.active = false; - a.header_remaining = 0; - a.pending_discontinuity = true; + a.drop_partial(); } return; } @@ -380,6 +460,12 @@ impl TsDemuxer { if idx < 0 { return; } + // transport_error_indicator: the packet is damaged, so its payload and CC are not + // trusted. Drop the open PES like a continuity break; the next PUSI resyncs. + if ts[1] & 0x80 != 0 { + self.assemblers[idx as usize].drop_partial(); + return; + } // adaptation_field_control == 0b00 is reserved (ISO 13818-1) and // carries no payload; discard so a corrupt/desynced packet can't // inject its 184 bytes into the PES assembler. @@ -416,41 +502,35 @@ impl TsDemuxer { // adaptation == 0x02 (AF only) already returned above, so only 0x03 // (AF + payload) can carry an adaptation field here. let discontinuity_flag = adaptation == 0x03 && ts[4] > 0 && (ts[5] & 0x80) != 0; - let cc_gap = cc_is_gap(asm.last_cc, cc); + // A duplicate repeats the previous packet (same CC and payload; the adaptation field, + // e.g. PCR, may differ): nothing new. A same-CC packet with other data is a gap. + let same_cc = asm.last_cc == Some(cc) && !discontinuity_flag; + if same_cc && asm.last_pusi == pusi && asm.last_payload == payload { + return; + } + let cc_gap = same_cc || cc_is_gap(asm.last_cc, cc); asm.last_cc = Some(cc); + asm.last_pusi = pusi; + asm.last_payload.clear(); + asm.last_payload.extend_from_slice(payload); // A gap means packets for this PID were lost. Sticky flag rides to the first // post-gap PES so the codec consumer drops forward to the next keyframe (B1). let gap = discontinuity_flag || cc_gap; if pusi { - // `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); - // Flush the previous PES first — a gap on this PUSI packet belongs to - // the PES starting now, not the one completing. Set `pending_discontinuity` - // after start(), else it wrongly stamps the pre-gap frame instead. - if let Some(prev) = asm.start(pts, dts, source) { + // Flush the previous PES first — a gap on this PUSI packet belongs to the PES + // starting now (as does a lost PES whose header never completed). Set + // `pending_discontinuity` after start(), else it stamps the pre-gap frame. + let lost_head = !asm.head.is_empty(); + if let Some(prev) = asm.start(None, None, source) { completed.push(prev); } - if gap { + if gap || lost_head { asm.pending_discontinuity = true; } - 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(); - } + asm.head.clear(); + asm.header_remaining = 0; + asm.take_head(payload); } else { // Non-PUSI continuation. if gap { @@ -463,13 +543,14 @@ impl TsDemuxer { pid = asm.pid, "TS continuity break on non-PUSI continuation; dropping partial PES", ); - asm.buffer.clear(); - asm.active = false; - asm.header_remaining = 0; + asm.drop_partial(); return; } } - if asm.header_remaining > 0 { + if !asm.head.is_empty() { + // The PES header is still arriving. + asm.take_head(payload); + } 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()); @@ -495,6 +576,32 @@ impl TsDemuxer { } } +// Stream ids whose PES has no header extension (ISO 13818-1 Table 2-22: program_stream_map, +// padding, private_stream_2, ECM, EMM, DSMCC, H.222.1 type E, program_stream_directory). +fn no_pes_extension(stream_id: u8) -> bool { + matches!( + stream_id, + 0xBC | 0xBE | 0xBF | 0xF0 | 0xF1 | 0xF2 | 0xF8 | 0xFF + ) +} + +// `parse_pes_header` of `data` once it holds the fixed header and any flagged PTS/DTS; +// `None` while more bytes are needed (and what arrived still matches a start code). +fn pes_header_complete(data: &[u8]) -> Option<(Option, Option, usize)> { + let need = match data { + [_, _, _, id, _, _, _, flags, hdl, ..] if !no_pes_extension(*id) => { + let (f, hdl) = (flags >> 6, *hdl as usize); + 9 + if f >= 2 && hdl >= 5 { 5 } else { 0 } + if f == 3 && hdl >= 10 { 5 } else { 0 } + } + _ => 9, + }; + let prefix = data.len().min(3); + if data.len() < need && data[..prefix] == [0, 0, 1][..prefix] { + return None; + } + Some(parse_pes_header(data)) +} + // Parse a PES header, extracting PTS/DTS. Returns `(pts, dts, header_len)` where // `header_len` is the FULL uncapped header length (9 + data_length, or 6 without // the extension; 0 = not a valid PES start) — caller skips it, carrying remainder. @@ -506,18 +613,7 @@ fn parse_pes_header(data: &[u8]) -> (Option, Option, usize) { let stream_id = data[3]; - // Some stream IDs don't carry the standard PES header extension - // (ISO 13818-1 Table 2-22: program_stream_map, padding, private_stream_2, - // ECM, EMM, DSMCC_stream 0xF2, H.222.1 type E 0xF8, program_stream_directory). - if stream_id == 0xBC - || stream_id == 0xBE - || stream_id == 0xBF - || stream_id == 0xF0 - || stream_id == 0xF1 - || stream_id == 0xF2 - || stream_id == 0xF8 - || stream_id == 0xFF - { + if no_pes_extension(stream_id) { return (None, None, 6); } @@ -526,9 +622,8 @@ fn parse_pes_header(data: &[u8]) -> (Option, Option, usize) { // this function and nothing shrinks `data` since, so no re-check here. let pts_dts_flags = (data[7] >> 6) & 0x03; let header_data_len = data[8] as usize; - // 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. + // Full, uncapped header length. `pes_header_complete` gathers the first 19 bytes (PTS/DTS) + // before parsing; only the *skip* length may extend into later packets. let header_len = 9 + header_data_len; let mut pts = None; @@ -611,9 +706,10 @@ fn ts_payload_base(pkt: &[u8]) -> Option { } } -// MPEG audio layer (1-3) from the first frame header at the start of `pid`'s -// first PES in `data`; `None` when no such PES or header is in the scan window. -fn mpeg_audio_layer(data: &[u8], pid: u16) -> Option { +// Per PID, the MPEG audio layer (1-3) from the frame header at the start of its first PES +// in `data` (`None` when that PES has no such header), found in one walk over `data`. +fn mpeg_audio_layers(data: &[u8]) -> std::collections::HashMap> { + let mut layers = std::collections::HashMap::new(); let mut offset = 0; while offset + BD_SOURCE_PACKET_BYTES <= data.len() { if !is_resync_point(data, offset) { @@ -622,30 +718,39 @@ fn mpeg_audio_layer(data: &[u8], pid: u16) -> Option { } let pkt = &data[offset..offset + BD_SOURCE_PACKET_BYTES]; let pkt_pid = (((pkt[5] & 0x1F) as u16) << 8) | pkt[6] as u16; - if pkt_pid == pid && pkt[5] & 0x40 != 0 { - let pes = &pkt[ts_payload_base(pkt)?..]; - if pes.get(..3)? != [0, 0, 1] { - return None; - } - let es = pes.get(9 + *pes.get(8)? as usize..)?; - let (b0, b1) = (*es.first()?, *es.get(1)?); - if b0 != 0xFF || b1 & 0xE0 != 0xE0 { - return None; - } - return match (b1 >> 1) & 0x03 { - 0 => None, - l => Some(4 - l), - }; + if pkt[5] & 0x40 != 0 { + layers + .entry(pkt_pid) + .or_insert_with(|| first_pes_layer(pkt)); } offset += BD_SOURCE_PACKET_BYTES; } - None + layers +} + +// MPEG audio layer of the frame header opening the PES that starts in `pkt`. +fn first_pes_layer(pkt: &[u8]) -> Option { + let pes = &pkt[ts_payload_base(pkt)?..]; + if pes.get(..3)? != [0, 0, 1] { + return None; + } + let es = pes.get(9 + *pes.get(8)? as usize..)?; + let (b0, b1) = (*es.first()?, *es.get(1)?); + if b0 != 0xFF || b1 & 0xE0 != 0xE0 { + return None; + } + match (b1 >> 1) & 0x03 { + 0 => None, + l => Some(4 - l), + } } // Reassemble a single PSI section (PAT/PMT) for `target_pid`/`table_id` across TS-packet // boundaries. Returns the section bytes (from table_id) or None. fn collect_psi_section(data: &[u8], target_pid: u16, table_id: u8) -> Option> { let mut offset = 0; + // First complete copy, used only if no copy passes `psi_section_ok`. + let mut fallback = None; while offset + BD_SOURCE_PACKET_BYTES <= data.len() { if !is_resync_point(data, offset) { offset += 1; @@ -679,7 +784,12 @@ fn collect_psi_section(data: &[u8], target_pid: u16, table_id: u8) -> Option= total { section.truncate(total); - return Some(section); + if psi_section_ok(§ion) { + return Some(section); + } + fallback.get_or_insert(section); + offset += BD_SOURCE_PACKET_BYTES; + continue; } // Need continuation packets: same PID, no PUSI. Use the canonical `cc_is_gap` // (§2.4.3.3) shared with `process_packet`, not a local test — a prior local @@ -729,14 +839,25 @@ fn collect_psi_section(data: &[u8], target_pid: u16, table_id: u8) -> Option= total { section.truncate(total); - return Some(section); + if psi_section_ok(§ion) { + return Some(section); + } + fallback.get_or_insert(section); + offset += BD_SOURCE_PACKET_BYTES; + continue; } // Incomplete section (truncated input) — stop looking. - return None; + return fallback; } offset += BD_SOURCE_PACKET_BYTES; } - None + fallback +} + +// A long-form PSI section that is current (current_next_indicator 1) and passes its +// CRC_32 (H.222.0 Annex A: over the whole section including the CRC, the register ends at 0). +fn psi_section_ok(section: &[u8]) -> bool { + section.len() >= 12 && section[5] & 0x01 == 1 && super::mpg::pack::crc32(section) == 0 } /// Scan BD-TS data for streams by parsing PAT and PMT tables. @@ -788,6 +909,7 @@ pub fn scan_streams(data: &[u8]) -> Option> { let prog_info_len = ((((pmt[10] & 0x0F) as usize) << 8) | pmt[11] as usize).min(end.saturating_sub(12)); let mut pos = 12 + prog_info_len; + let mut layers = None; while pos + 5 <= end { let stream_type = pmt[pos]; @@ -798,7 +920,14 @@ pub fn scan_streams(data: &[u8]) -> Option> { // absent from Blu-ray's STN table. Keep that distinction local to TS. let codec = match stream_type { // MPEG-1/2 audio covers Layers I-III; only the ES says which. - 0x03 | 0x04 if mpeg_audio_layer(data, es_pid) == Some(3) => Codec::Mp3, + 0x03 | 0x04 + if layers + .get_or_insert_with(|| mpeg_audio_layers(data)) + .get(&es_pid) + == Some(&Some(3)) => + { + Codec::Mp3 + } 0x03 | 0x04 => Codec::Mp2, 0x0f => Codec::Aac, 0x87 => Codec::Ac3Plus, @@ -973,6 +1102,156 @@ mod tests { v } + // A duplicate packet (same CC, h222 2.4.3.3) carries no new data: neither a repeated + // continuation nor a repeated PUSI may change the PES. + #[test] + fn duplicate_packets_are_ignored() { + let pid = 0x1011; + let mut demux = TsDemuxer::new(&[pid]); + let start = ts_payload_packet(pid, true, 0, &pes_start(&[b'A'; 175])); + let cont = ts_payload_packet(pid, false, 1, &[b'B'; 184]); + let mut out = Vec::new(); + for p in [&start, &start, &cont, &cont] { + out.extend(demux.feed(p)); + } + out.extend(demux.flush()); + assert_eq!(out.len(), 1, "one PES"); + let mut want = vec![b'A'; 175]; + want.extend_from_slice(&[b'B'; 184]); + assert_eq!(out[0].data, want); + assert!(!out[0].discontinuity); + } + + // h222 lets a PES header span packets: a PUSI payload too short for the fixed header or + // its PTS/DTS is completed by the continuation, and the timestamps are kept. + #[test] + fn pes_header_split_across_packets_keeps_its_timestamps() { + let pid = 0x1011; + // PTS 0x1_2345_6789 and DTS 0x1_2345_0000 ('11' flags, header_data_length 10). + let mut hdr = vec![0x00, 0x00, 0x01, 0xE0, 0x00, 0x00, 0x80, 0xC0, 0x0A]; + hdr.extend_from_slice(&[0x39, 0x8D, 0x15, 0xCF, 0x13]); + hdr.extend_from_slice(&[0x19, 0x8D, 0x15, 0x00, 0x01]); + for split in [4, 11, 16] { + let mut demux = TsDemuxer::new(&[pid]); + let mut rest = hdr[split..].to_vec(); + rest.extend_from_slice(b"ES"); + let mut out = demux.feed(&es_packet_exact(pid, true, &hdr[..split])); + out.extend(demux.feed(&es_packet_exact(pid, false, &rest))); + out.extend(demux.flush()); + assert_eq!(out.len(), 1, "split {split}"); + assert_eq!(out[0].data, b"ES", "split {split}"); + assert_eq!(out[0].pts, Some(0x1_2345_6789), "split {split}"); + assert_eq!(out[0].dts, Some(0x1_2345_0000), "split {split}"); + } + } + + // A byte slip mid-stream loses only the packets it hits: the demuxer resyncs. + #[test] + fn byte_slip_resyncs_to_the_next_packet_boundary() { + let pid = 0x1011; + let mut demux = TsDemuxer::new(&[pid]); + let mut data = ts_payload_packet(pid, true, 0, &pes_start(b"AAAA")); + data.push(0x00); + data.extend(ts_payload_packet(pid, false, 1, b"BBBB")); + data.extend(ts_payload_packet(pid, true, 2, &pes_start(b"CCCC"))); + let mut out = demux.feed(&data); + out.extend(demux.flush()); + assert_eq!(out.len(), 2, "{out:?}"); + assert!(out[0].data.starts_with(b"AAAA")); + assert!(out[0].data[175..].starts_with(b"BBBB"), "continuation kept"); + assert!(out[1].data.starts_with(b"CCCC")); + } + + // A lone damaged sync byte keeps the grid even when stray 0x47 bytes 192 apart would + // pass for an off-grid packet boundary. + #[test] + fn resync_prefers_the_next_grid_slot_over_an_off_grid_pair() { + let pid = 0x1011; + let mut demux = TsDemuxer::new(&[pid]); + let mut bad = ts_payload_packet(pid, false, 1, &[0xAA; 184]); + bad[4] = 0x00; + bad[50] = SYNC_BYTE; + let mut b = ts_payload_packet(pid, true, 2, &pes_start(&[0xBB; 60])); + b[50] = SYNC_BYTE; + let mut data = ts_payload_packet(pid, true, 0, &pes_start(b"AAAA")); + data.extend(bad); + data.extend(b); + let mut out = demux.feed(&data); + out.extend(demux.flush()); + assert_eq!(out.len(), 2, "{out:?}"); + assert!(out[1].data.starts_with(&[0xBB; 30])); + } + + // Zero-filled gaps (unreadable sectors) keep the packet grid: a lone packet after one + // is still read, without a second packet to corroborate it. + #[test] + fn lone_packets_between_zero_filled_gaps_are_kept() { + let pid = 0x1011; + let mut demux = TsDemuxer::new(&[pid]); + let mut data = vec![0u8; 5 * BD_SOURCE_PACKET_BYTES]; + data.extend(ts_payload_packet(pid, true, 0, &pes_start(b"AAAA"))); + data.extend(vec![0u8; 3 * BD_SOURCE_PACKET_BYTES]); + data.extend(ts_payload_packet(pid, true, 1, &pes_start(b"BBBB"))); + data.extend(vec![0u8; 2 * BD_SOURCE_PACKET_BYTES]); + let mut out = demux.feed(&data); + out.extend(demux.flush()); + assert_eq!(out.len(), 2, "{out:?}"); + assert!(out[1].data.starts_with(b"BBBB")); + } + + // A packet flagged transport_error_indicator is damaged: its PES is dropped, not + // spliced, and the next PES is flagged. + #[test] + fn transport_error_packet_drops_its_pes() { + let pid = 0x1011; + let mut demux = TsDemuxer::new(&[pid]); + let mut bad = ts_payload_packet(pid, false, 1, b"XXXX"); + bad[5] |= 0x80; + let mut out = demux.feed(&ts_payload_packet(pid, true, 0, &pes_start(b"AAAA"))); + out.extend(demux.feed(&bad)); + out.extend(demux.feed(&ts_payload_packet(pid, true, 2, &pes_start(b"CCCC")))); + out.extend(demux.flush()); + assert_eq!(out.len(), 1, "A is dropped: {out:?}"); + assert!(out[0].data.starts_with(b"CCCC")); + assert!(out[0].discontinuity); + } + + // Clips joined through one demuxer restart the CC arbitrarily: a new clip's PUSI with + // the previous packet's CC is new data, not a duplicate, and must start its own PES. + #[test] + fn same_cc_packet_with_new_data_is_not_a_duplicate() { + let pid = 0x1011; + let mut demux = TsDemuxer::new(&[pid]); + let mut out = Vec::new(); + for p in [ + ts_payload_packet(pid, true, 0, &pes_start(b"AAAA")), + ts_payload_packet(pid, false, 1, &[b'a'; 184]), + ts_payload_packet(pid, true, 1, &pes_start(b"BBBB")), + ts_payload_packet(pid, false, 2, &[b'b'; 184]), + ] { + out.extend(demux.feed(&p)); + } + out.extend(demux.flush()); + assert_eq!(out.len(), 2, "{out:?}"); + assert!(out[0].data.starts_with(b"AAAA") && out[0].data.ends_with(&[b'a'; 184])); + assert!(out[1].data.starts_with(b"BBBB") && out[1].data.ends_with(&[b'b'; 184])); + } + + // A PUSI whose payload is not a PES start leaves no PES open: the continuations after it + // are not emitted as a headless frame, and the next PES is flagged. + #[test] + fn pusi_without_a_pes_header_opens_no_pes() { + let pid = 0x1011; + let mut demux = TsDemuxer::new(&[pid]); + let mut out = demux.feed(&ts_payload_packet(pid, true, 0, &[0xAB; 184])); + out.extend(demux.feed(&ts_payload_packet(pid, false, 1, &[0xCD; 184]))); + out.extend(demux.feed(&ts_payload_packet(pid, true, 2, &pes_start(b"NEXT")))); + out.extend(demux.flush()); + assert_eq!(out.len(), 1, "only the valid PES: {out:?}"); + assert!(out[0].data.starts_with(b"NEXT")); + assert!(out[0].discontinuity, "the dropped unit is a gap"); + } + // Regression: a non-PUSI continuation with a CC gap means packets were // dropped — the partial PES must be discarded, not spliced onto. A clean // follow-on PUSI then produces exactly the next PES, not a corrupt splice. @@ -1228,6 +1507,19 @@ mod tests { // ── scan_streams PMT parsing ────────────────────────────────────────── + // Next continuity_counter for `pid` in this test thread, so fixture packets form a + // well-formed per-PID sequence (a repeated CC is a duplicate packet). + fn next_cc(pid: u16) -> u8 { + thread_local!(static CC: std::cell::RefCell> = + Default::default()); + CC.with(|m| { + let mut m = m.borrow_mut(); + let c = m.entry(pid).or_insert(0x0F); + *c = (*c + 1) & 0x0F; + *c + }) + } + /// Wrap a 188-byte TS packet body in a 192-byte BD-TS packet /// (4-byte timecode prefix the scanner skips). fn bdts_packet(body: [u8; 184], pid: u16, pusi: bool) -> Vec { @@ -1239,7 +1531,7 @@ mod tests { pkt[5] |= 0x40; } pkt[6] = (pid & 0xFF) as u8; - pkt[7] = 0x10; // payload only, no adaptation field + pkt[7] = 0x10 | next_cc(pid); // payload only, no adaptation field pkt[8..8 + 184].copy_from_slice(&body); pkt } @@ -1266,7 +1558,7 @@ mod tests { body[i + 9] = 0x01; body[i + 10] = 0xE0 | (((pmt_pid >> 8) as u8) & 0x1F); body[i + 11] = (pmt_pid & 0xFF) as u8; - // (CRC bytes left as 0xFF — scanner doesn't validate CRC) + // (CRC left as 0xFF: with no valid copy the scanner uses the first one) let _ = &mut i; bdts_packet(body, 0, true) } @@ -1314,7 +1606,7 @@ mod tests { pkt[5] |= 0x40; } pkt[6] = (pid & 0xFF) as u8; - pkt[7] = 0x10; // payload only, no adaptation field + pkt[7] = 0x10 | next_cc(pid); // payload only, no adaptation field let room = TS_PACKET_BYTES - 4; // 184 ES bytes after the 4-byte TS header let n = payload.len().min(room); pkt[8..8 + n].copy_from_slice(&payload[..n]); @@ -1386,30 +1678,11 @@ mod tests { "garbage PUSI packet must not complete a PES on its own" ); - // Continuation packet (no PUSI) carrying real ES bytes. - let es = [0xDEu8, 0xAD, 0xBE, 0xEF]; - stream.extend(demux.feed(&data_packet(pid, false, &es))); + // A continuation after it belongs to no PES (headless): nothing is emitted, so + // none of the garbage (0xAA, the embedded 00 00 01) can reach the ES. + stream.extend(demux.feed(&data_packet(pid, false, &[0xDE, 0xAD, 0xBE, 0xEF]))); stream.extend(demux.flush()); - - assert_eq!(stream.len(), 1, "one PES assembled from the continuation"); - let pes = &stream[0]; - // The continuation ES bytes survive… - assert!( - pes.data.windows(es.len()).any(|w| w == es), - "continuation ES bytes present, got {:02X?}", - pes.data - ); - // …but none of the garbage PUSI payload leaked in. In particular the - // 0xAA filler and the embedded 00 00 01 sequence must be absent — the - // malformed PES header contributed ZERO bytes to the elementary stream. - assert!( - !pes.data.contains(&0xAA), - "garbage PES-header bytes must not appear in the elementary stream" - ); - assert!( - !pes.data.windows(3).any(|w| w == [0x00, 0x00, 0x01]), - "no injected start code leaked from the malformed PES header" - ); + assert!(stream.is_empty(), "no PES without a PES header: {stream:?}"); } #[test] @@ -1432,6 +1705,51 @@ mod tests { ); } + // Many MPEG-audio PMT entries over a sync-free head: one walk, each PID keeps its layer. + #[test] + fn scan_streams_labels_many_mpeg_audio_pids_in_one_walk() { + use crate::disc::{Codec, Stream}; + let mut entries = vec![(0x1b, 0x1011)]; + entries.extend((0..60u16).map(|i| (0x03, 0x1100 + i))); + let mut data = pat_packet(0x100); + data.extend(pmt_two_packets(0x100, &entries)); + for (pid, b1) in [(0x1100, 0xFB), (0x1101, 0xFD)] { + let mut pes = vec![0x00, 0x00, 0x01, 0xC0, 0x00, 0x00, 0x80, 0x00, 0x00]; + pes.extend_from_slice(&[0xFF, b1, 0x90, 0x64]); + data.extend(es_packet_exact(pid, true, &pes)); + } + data.resize(data.len() + (1 << 20), 0); + let streams = scan_streams(&data).unwrap(); + assert!(matches!(&streams[1], Stream::Audio(a) if a.codec == Codec::Mp3)); + assert!(matches!(&streams[2], Stream::Audio(a) if a.codec == Codec::Mp2)); + } + + // Give a single-packet PSI fixture (pointer_field 0) its real CRC_32. + fn with_psi_crc(mut pkt: Vec) -> Vec { + let total = 3 + ((((pkt[10] & 0x0F) as usize) << 8) | pkt[11] as usize); + let crc = crate::mux::mpg::pack::crc32(&pkt[9..9 + total - 4]); + pkt[9 + total - 4..9 + total].copy_from_slice(&crc.to_be_bytes()); + pkt + } + + // A PMT copy failing its CRC_32 (or not current) is skipped for a later valid copy. + #[test] + fn scan_streams_prefers_a_crc_valid_current_pmt() { + use crate::disc::{Codec, Stream}; + let good = [(0x1b, 0x1011), (0x81, 0x1100)]; + let mut flipped = with_psi_crc(pmt_packet(0x100, &good)); + flipped[9 + 12 + 5] = 0x87; // second entry's stream_type, after the CRC was set + let mut next = with_psi_crc(pmt_packet(0x100, &[(0x1b, 0x1011), (0x87, 0x1100)])); + next[9 + 5] &= !0x01; // current_next_indicator 0: not yet applicable + let next = with_psi_crc(next); + let mut data = with_psi_crc(pat_packet(0x100)); + for p in [flipped, next, with_psi_crc(pmt_packet(0x100, &good))] { + data.extend(p); + } + let streams = scan_streams(&data).unwrap(); + assert!(matches!(&streams[1], Stream::Audio(a) if a.codec == Codec::Ac3)); + } + #[test] fn scan_streams_maps_standard_transport_audio_types() { use crate::disc::{Codec, Stream}; @@ -1668,7 +1986,7 @@ mod tests { } // A duplicate TS packet (same CC, identical payload) is explicitly legal; - // `process_packet` already tolerates it. The PSI reassembler must too: + // `process_packet` skips it. The PSI reassembler must too: // treat it as a duplicate (not appended twice), not as a desync. #[test] fn scan_streams_tolerates_duplicate_pmt_continuation_packet() { @@ -1765,6 +2083,7 @@ mod tests { let payload_off = 8 + pad; pkt[payload_off..payload_off + payload.len()].copy_from_slice(payload); } + pkt[7] |= next_cc(pid); pkt } @@ -2381,10 +2700,9 @@ mod tests { pes_start.extend_from_slice(&[0xAB; 10]); demux.feed(&es_packet_exact(pid, true, &pes_start)); let payload = [0xCCu8; 184]; - let cont_pkt = data_packet(pid, false, &payload); let mut high_water = 0usize; for _ in 0..(expected_share / 184 + 64) { - demux.feed(&cont_pkt); + demux.feed(&data_packet(pid, false, &payload)); let idx = demux.pid_index[pid as usize] as usize; high_water = high_water.max(demux.assemblers[idx].buffer.len()); } @@ -2399,6 +2717,27 @@ mod tests { ); } + // The buffer's allocation, not just its length, stays within the PID's share: doubling + // past a non-power-of-two cap would overshoot the aggregate ceiling. + #[test] + fn pes_buffer_capacity_stays_within_the_pid_share() { + let pids: Vec = (0..48).map(|i| 0x1100 + i).collect(); + let mut demux = TsDemuxer::new(&pids); + let pid = pids[0]; + let cap = demux.assemblers[0].cap; + assert!(!cap.is_power_of_two()); + demux.feed(&es_packet_exact(pid, true, &pes_start(&[]))); + for _ in 0..cap / 184 { + demux.feed(&data_packet(pid, false, &[0xCC; 184])); + let a = &demux.assemblers[0]; + assert!( + a.buffer.capacity() <= cap, + "{} > {cap}", + a.buffer.capacity() + ); + } + } + #[test] fn pes_buffer_cap_resets_on_overflow_and_recovers_on_next_pusi() { // Feed continuation-only packets past MAX_PES_BUFFER, then verify the buffer @@ -2415,11 +2754,10 @@ mod tests { // Continuation packets with 184-byte payloads, no PUSI. Each call to // feed() processes one 192-byte BD-TS packet. let payload = [0xCCu8; 184]; - let cont_pkt = data_packet(pid, false, &payload); let packets_needed = MAX_PES_BUFFER / 184 + 2; let mut mid_out: Vec = Vec::new(); for _ in 0..packets_needed { - mid_out.extend(demux.feed(&cont_pkt)); + mid_out.extend(demux.feed(&data_packet(pid, false, &payload))); // Verify the internal buffer is bounded: no assembler may hold // more than MAX_PES_BUFFER bytes at any point. for asm in &demux.assemblers { diff --git a/src/mux/tsmux.rs b/src/mux/tsmux.rs index 18269e2a..9c682a5b 100644 --- a/src/mux/tsmux.rs +++ b/src/mux/tsmux.rs @@ -19,14 +19,14 @@ use crate::consts::TS_PAYLOAD_BYTES; /// 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 = 0x1011..=0x101F; +pub(crate) const VIDEO_PID_RANGE: std::ops::RangeInclusive = 0x1011..=0x101F; // Largest PES payload that fits a bounded `PES_packet_length` (u16) on a // `0xBD` stream after the 8 PES-header bytes. Frames larger than this split // into multiple PES (unbounded `0` length is video-only). const MAX_BD_PES_PAYLOAD: usize = u16::MAX as usize - 8; -fn is_video_pid(pid: u16) -> bool { +pub(crate) fn is_video_pid(pid: u16) -> bool { VIDEO_PID_RANGE.contains(&pid) } @@ -807,6 +807,8 @@ mod tests { // IDR: should carry codec_private NALs prepended. let idr = fake_hevc_nal(19, 50); mux.write_frame(0, 41_000_000, true, &idr).unwrap(); + // A later IDR: the parameter sets are not prepended again. + mux.write_frame(0, 82_000_000, true, &idr).unwrap(); mux.finish().unwrap(); } let packets = parse_bd_ts(&sink); @@ -831,6 +833,12 @@ mod tests { pos_marker < pos_idr, "codec_private must precede IDR in TS payload" ); + let markers = video_bytes.windows(marker.len()).filter(|w| *w == marker); + assert_eq!( + markers.count(), + 1, + "codec_private only before the first IDR" + ); } #[test] @@ -1084,6 +1092,34 @@ mod tests { /// The default (`video_codec` = HEVC) still converts, so the test above is /// pinning the flag rather than a no-op. Same input, opposite expectation. + // An avcC declaring 2-octet NAL lengths (lengthSizeMinusOne 1) must be converted with + // 2-octet prefixes, not the 4-octet default. + #[test] + fn avcc_nal_length_size_drives_the_annex_b_conversion() { + let (a, b): (&[u8], &[u8]) = (&[0x65, 0x88, 0x84, 0x21, 0x43], &[0x06, 0x05, 0x01]); + let mut es = Vec::new(); + for nal in [a, b] { + es.extend_from_slice(&(nal.len() as u16).to_be_bytes()); + es.extend_from_slice(nal); + } + let mut sink: Vec = Vec::new(); + { + let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]); + mux.set_video_codec(0, crate::disc::Codec::H264).unwrap(); + // avcC header only: lengthSizeMinusOne = 1, no SPS/PPS. + mux.set_codec_private(0, vec![1, 0x42, 0xC0, 0x1E, 0xFD, 0xE0, 0]) + .unwrap(); + mux.write_frame(0, 0, true, &es).unwrap(); + mux.finish().unwrap(); + } + let out = reassemble_es(&parse_bd_ts(&sink), VIDEO_PID); + let mut want = vec![0, 0, 0, 1]; + want.extend_from_slice(a); + want.extend_from_slice(&[0, 0, 0, 1]); + want.extend_from_slice(b); + assert_eq!(out, want); + } + #[test] fn nal_video_es_is_converted_to_annex_b_by_default() { let es: Vec = vec![0x00, 0x00, 0x00, 0x06, 0xB3, 0x12, 0x34, 0x56, 0x78, 0x9A]; @@ -1434,7 +1470,27 @@ mod tests { /// mutated definition would still pass those assertions. #[test] fn max_bd_pes_payload_has_the_documented_value() { - assert_eq!(MAX_BD_PES_PAYLOAD, u16::MAX as usize - 8); + assert_eq!(MAX_BD_PES_PAYLOAD, 65_527); + } + + // A private_stream_1 frame of exactly the limit fills one PES to PES_packet_length + // 0xFFFF; one byte more splits into two. + #[test] + fn bd_audio_split_boundary_is_exact() { + for (len, pes) in [(65_527usize, 1usize), (65_528, 2)] { + let mut sink: Vec = Vec::new(); + { + let mut mux = TsMuxer::new(&mut sink, &[AUDIO_PID]); + mux.write_frame(0, 0, false, &vec![0x5A; len]).unwrap(); + mux.finish().unwrap(); + } + let packets = parse_bd_ts(&sink); + let starts: Vec<&TsPacket> = packets.iter().filter(|p| p.pusi).collect(); + assert_eq!(starts.len(), pes, "{len} bytes"); + let first = u16::from_be_bytes([starts[0].payload[4], starts[0].payload[5]]); + assert_eq!(first, 0xFFFF, "{len} bytes: the first PES is full"); + assert_eq!(reassemble_es(&packets, AUDIO_PID), vec![0x5A; len]); + } } // An oversized video access unit must stay ONE PES (unbounded-length